diff --git a/configs/nemotron_duplex.yaml b/configs/nemotron_duplex.yaml index 35441b5b79..c49e0b8e24 100644 --- a/configs/nemotron_duplex.yaml +++ b/configs/nemotron_duplex.yaml @@ -9,11 +9,19 @@ model: "nemotron_duplex" max_seq_len: 12288 # config: voicechat streaming_max_len / vllm_max_model_len -# One recurrent slot per session (DEFAULT_MAMBA_SLOTS); past it the engine spins -# on admit and wedges every session. Keep the two equal; a live call past the -# cap gets a 503, and only one-shot requests queue. +# Sessions at once; a live call past it gets a 503 (only one-shot requests +# queue). Each session holds a slot in +# resources mamba_state, encoder_state, encoder_front and encoder_kv and its +# pages in nano_kv and talker_kv, all sized for 32 (DEFAULT_SESSIONS); to raise +# the cap, raise their max_slots (cap + 1 for the recurrent pools, for the +# sink; cap for encoder_kv) and the KV pools' max_num_pages with it. +# The server refuses to start with a cap one of them cannot hold. max_concurrent_requests: 32 +# No per-message deflate on the WebSocket routes: compressing 32 live calls' +# 80 ms appends cost the event loop ~8% on /generate/ws (H100, 2026-10-17). +ws_per_message_deflate: false + # A live call's input files over its whole life; past it the call ends with a # 400. 160 s of 16 kHz audio is about 10 MB as float32 WAV: 32 leaves room for # stereo or a 48 kHz opening clip. A file is decoded at most one frame past diff --git a/configs/nemotron_duplex_disagg.yaml b/configs/nemotron_duplex_disagg.yaml index 61cfebd1d3..4b146a19c7 100644 --- a/configs/nemotron_duplex_disagg.yaml +++ b/configs/nemotron_duplex_disagg.yaml @@ -14,11 +14,19 @@ model: "nemotron_duplex" max_seq_len: 12288 -# One recurrent slot per session (DEFAULT_MAMBA_SLOTS); past it the engine spins -# on admit and wedges every session. Keep the two equal; a live call past the -# cap gets a 503, and only one-shot requests queue. +# Sessions at once; a live call past it gets a 503 (only one-shot requests +# queue). Each session holds a slot in +# resources mamba_state, encoder_state, encoder_front and encoder_kv and its +# pages in nano_kv and talker_kv, all sized for 32 (DEFAULT_SESSIONS); to raise +# the cap, raise their max_slots (cap + 1 for the recurrent pools, for the +# sink; cap for encoder_kv) and the KV pools' max_num_pages with it. +# The server refuses to start with a cap one of them cannot hold. max_concurrent_requests: 32 +# No per-message deflate on the WebSocket routes: compressing 32 live calls' +# 80 ms appends cost the event loop ~8% on /generate/ws (H100, 2026-10-17). +ws_per_message_deflate: false + # A live call's input files over its whole life; past it the call ends with a # 400. 160 s of 16 kHz audio is about 10 MB as float32 WAV: 32 leaves room for # stereo or a 48 kHz opening clip. A file is decoded at most one frame past diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 45c884f69b..2e5096f479 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -569,9 +569,12 @@ The spec types are: ``SamplingReqConfig`` above. * - ``RecurrentStateSpec(config=RecurrentStateConfig(...))`` - Fixed-size per-request state that carries between steps (an LSTM's ``h`` and - ``c``, a cursor, a conv tail): one slot per request, ``blocks`` naming a shape and - dtype for each piece, ``num_layers`` copies of each, ``max_slots`` slots with slot 0 - the sink that padding rows address. ``declare_step`` gives each row a + ``c``, a cursor, a conv tail, a Mamba or gated-DeltaNet layer's conv window and SSM + state, a streaming encoder's caches): one slot per request and label, ``blocks`` + naming a shape and dtype for each piece (dtype None: the engine's compute dtype, + what its KV caches store), ``num_layers`` copies of each, ``max_slots`` slots + (deployment-tunable) with slot 0 the sink that padding rows address. A recurrent + kernel reads the step's ``addressing``. ``declare_step`` gives each row a ``RecurrentStep`` segment with span >= 1. The forward reads its rows with ``pool.gather(name, start=...)`` and writes them back with ``pool.scatter_(name, value)``; both follow the step's plan, so there is no @@ -866,7 +869,9 @@ within itself and matches the offline mask. ``rel_bias`` is then indexed by that ``head_dim ** -0.5``). A per-query window takes no sink, and its right context is 0. ``retention`` may be a single ``SinkWindow`` when it does not depend on the source length. ``BoundedKVConfig(dtype=None)`` stores the cache in the engine's KV dtype, for layers that -compute in it. +compute in it. Nemotron VoiceChat's streaming encoder (``ENCODER_KV``, +``encoder_kv_config`` in ``mstar/model/nemotron_duplex/components/conformer.py``) is the +per-query example. Step 3 — Declare the computation graph -------------------------------------- diff --git a/docs/clients.rst b/docs/clients.rst index ac2ef4c515..357c970518 100644 --- a/docs/clients.rst +++ b/docs/clients.rst @@ -98,6 +98,9 @@ WebSocket ``/generate/ws`` ``/generate/ws`` is the same request over one persistent WebSocket, for control loops that cannot afford an HTTP round trip per step (robot policies, streaming world models). +The server negotiates per-message deflate on its WebSocket routes unless the deployment +turns it off (``ws_per_message_deflate`` in :doc:`serving`, as the Nemotron VoiceChat +configs do). Each message is one request with the ``/generate`` fields — ``text``, ``files`` as ``[{"name": ..., "data": ...}]``, ``input_modalities``, ``output_modalities``, ``model_kwargs``, ``request_id`` — sent either as a JSON text frame (``data`` base64) or @@ -189,6 +192,12 @@ is at ``max_concurrent_requests`` ends at once with ``status`` 503, rather than slot and start that far behind its input; a one-shot request waits in the queue. Other models ignore ``end_input``, and an append to them is rejected. Live input needs the Python frontend: ``--rust-frontend`` serves no ``/generate/ws``. +The append size that serves best is the model's: Nemotron VoiceChat runs a live step as one +CUDA graph only for an append of exactly 80 ms (1280 samples at 16 kHz) onto a call whose +audio so far is a whole number of 80 ms frames. Keep the opening clip and every append a +multiple of 1280 samples: one that is not (a one-second opening clip, 16000 samples) keeps +the call off the graph for the rest of its life. Other sizes are served the same, more +slowly (see :doc:`models`). .. code-block:: python diff --git a/docs/models.rst b/docs/models.rst index 478484fb12..d87f850788 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -597,7 +597,8 @@ Nemotron VoiceChat (``nemotron_duplex``) notes PCM: each 5-frame chunk is decoded with the previous 3 frames as left context, which covers the decoder's receptive field, and a window's last 12 samples, the inverse-STFT overlap, are held back and emitted with the next chunk, so - the stream equals the full-history decode sample for sample). + the stream equals the full-history decode sample for sample; that steady + window of 8 frames replays a CUDA graph, other windows decode eagerly). The decode loops are stream-terminated: they end when the upstream stream closes, not on EOS, which is an ordinary per-frame token in duplex speech. - ``nano_llm`` declares a paged KV cache over its four attention layers with @@ -605,10 +606,22 @@ Nemotron VoiceChat (``nemotron_duplex``) notes positional encoding, so no position resource), a recurrent-state pool holding the 27 Mamba-2 layers' conv window and SSM state, one slot per session (``mamba_state``: 32 sessions by default, ~137 MB per slot in fp32, tunable - with ``max_slots`` / ``state_dtype``; the shipped configs set - ``max_concurrent_requests: 32`` to match, so raise the two together), the Mamba-2 resource planned on it + with ``max_slots``), the Mamba-2 resource planned on it (``mamba``) and the agent-text sampler (``nano_sampler``). Its decode step is captured as a CUDA graph for batch sizes 1 to 32. +- Every per-session resource is sized for 32 sessions: a slot in each of + ``mamba_state``, ``encoder_state`` and ``encoder_front`` (``max_slots``, one + more than the sessions, for the sink) and in ``encoder_kv`` (``max_slots``, + no sink), and its pages in ``nano_kv`` and + ``talker_kv`` (``max_num_pages``; both sized from ``max_session_s``, and the + nano's from ``max_prompt_tokens`` too). The shipped configs set + ``max_concurrent_requests: 32`` to match; to serve more, raise the cap and + all six with it. The server + refuses to start with a cap one of them cannot hold, since a session past + any of them would wait on admit, and its frames with it, until another ends. + The nano, the talker, the encoder and the codec each run at most 32 sessions + a step, their widest CUDA graph; more are split over several steps rather + than run eagerly. - ``eartts_talker`` advances every live session in one backbone pass per frame through the engine's paged KV (``talker_kv`` / ``talker_attn`` / ``talker_pos``): each session owns two streams, the conditional and the @@ -634,8 +647,12 @@ Nemotron VoiceChat (``nemotron_duplex``) notes fast as it can). Live calls paced in real time over ``/generate/ws`` (``test/nemotron_duplex/live_ws_request.py``, 80 ms appends, 2026-10-17, one H100): input-to-token lag p50 48 ms with one call and about 100-120 ms with - 8, which keep up; from 24 calls up the server falls behind real time, so lag - grows over the call. Served audio is checked for intelligibility by + 8, which keep up; from 24 calls up the server fell behind real time, so lag + grew over the call. Capturing the encoder's live step (below) took its host + time at 32 calls from 23.6 to 8.4 ms and cut live lag 2-6x at 16-32 calls + (paced runs interleaved against the uncaptured encoder, ``MSTAR_SHM_ARENA=1``); + 32 calls staggered over 2 s all finish, heard latency p50 2.1-2.3 s (from the + user stopping to agent audio playable without gaps). Served audio is checked for intelligibility by transcribing it (``test/nemotron_duplex/asr_check.py``, Whisper large-v3-turbo): MaskGIT sampling is knife-edge, so attention-backend numerics change the waveform but not the words. @@ -682,6 +699,21 @@ Nemotron VoiceChat (``nemotron_duplex``) notes (inferred from those frames, not measured on served text); under the served bf16 a different chunking can flip a knife-edge frame such as the agent's turn start. +- The encoder keeps each call's state in a slot per session of three + resources: its attention keys and values in a bounded KV (``encoder_kv``), + which keeps a call's last 70 frames per layer and runs the attention, each new + frame over the 70 before it, with the relative-position term as a band of + distances; its conv inputs (``encoder_state``); and its newest audio and mel + frames (``encoder_front``), the last two recurrent-state pools. A + call that appends one frame of audio (1280 samples) on the frame grid, every + live step of an 80 ms client, runs its whole step as a CUDA graph (the grid is + the call's start: an opening clip or append that is not a multiple of 1280 + samples keeps the call off it for good); other calls take the front end + eagerly, then the layer stack as a graph when each got at most one frame (a + live step, the call's end), or eagerly (a one-shot upload, a large append: + a chunk over a second is not batched with other calls, so the engine runs + each call of that step on its own). Appends of other sizes are served the same, + without the whole-step graph. - Checks: ``test/nemotron_duplex/duplex_request.py`` prepares a clip from the bundled ``turn_taking.wav`` and sends it; ``parity_compare.py`` runs the standalone reference path on the same clip and checks the served text diff --git a/docs/serving.rst b/docs/serving.rst index 5ba14eebc1..65d3d2f4ab 100644 --- a/docs/serving.rst +++ b/docs/serving.rst @@ -273,6 +273,14 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key ``max_call_input_mb`` it counts what the model was handed, not bytes uploaded. * - ``model_kwargs`` - *(optional)* Server-init model parameters (see below). + * - ``ws_per_message_deflate`` + - *(optional, default true)* Whether the server offers per-message deflate on its + WebSocket routes (``/generate/ws``, ``/v1/realtime``); it applies to every route. + Deflate shrinks 80 ms speech frames to 49-73% of their size as base64 JSON and to + 55-90% as binary PCM, and digital silence to 1-2%, but compressing every frame + costs the event loop: about 8% of it at 32 live Nemotron VoiceChat calls on + ``/generate/ws``, whose configs turn it off. Turn it off where the server's + event loop, not the network, is the bottleneck. * - ``speculation_mixes_walks`` - *(optional)* Per node, ``true``/``false`` for every walk of the node's combined walks, or a mapping of walk -> ``true``/``false``. On for a walk, a speculation diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index 4b6513d308..c9fde9ca9b 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -2382,6 +2382,16 @@ async def shutdown_event(): # CLI entry point # ------------------------------------------------------------------ +def _uvicorn_config(app, host: str, port: int, config: dict) -> uvicorn.Config: + # Per-message deflate is negotiated server-wide, for every WebSocket route, + # so it is the deployment's to turn off (see ws_per_message_deflate in + # docs/serving.rst); on by default, as uvicorn's. + return uvicorn.Config( + app, host=host, port=port, access_log=False, + ws_per_message_deflate=bool(config.get("ws_per_message_deflate", True)), + ) + + def _set_preprocess_threads(model, config: dict) -> None: """Size torch's intra-op pool for this process's media preprocessing. @@ -2593,9 +2603,7 @@ def main(argv: list[str] | None = None): logger.info("Starting mstar API server on %s:%s", args.host, args.port) # uvicorn.run() inlined so the server object is reachable. The # message thread stops it when the conductor dies. - server = uvicorn.Server( - uvicorn.Config(app, host=args.host, port=args.port, access_log=False) - ) + server = uvicorn.Server(_uvicorn_config(app, args.host, args.port, config)) def _stop_server(): server.should_exit = True diff --git a/mstar/engine/resources/linear_attn/base.py b/mstar/engine/resources/linear_attn/base.py index 08822b7653..3ace685cf0 100644 --- a/mstar/engine/resources/linear_attn/base.py +++ b/mstar/engine/resources/linear_attn/base.py @@ -40,7 +40,7 @@ def build(cls, spec: LinearAttnSpec, info: EngineResourceInfo): config=spec.config, geometry=Mamba2Geometry.from_blocks(pool_config.blocks), num_layers=pool_config.num_layers, - state_dtype=pool_config.blocks["ssm"].dtype, + state_dtype=pool_config.blocks["ssm"].resolved_dtype(info.kv_dtype), has_sink=not pool_config.disable_sink_slot, device=info.device, ) @@ -55,7 +55,7 @@ def build(cls, spec: LinearAttnSpec, info: EngineResourceInfo): max_rows_in_combined=info.max_rows_in_combined, geometry=geometry, num_layers=pool_config.num_layers, - state_dtype=pool_config.blocks["state"].dtype, + state_dtype=pool_config.blocks["state"].resolved_dtype(info.kv_dtype), has_sink=not pool_config.disable_sink_slot, device=info.device, ) diff --git a/mstar/engine/resources/recurrent/config.py b/mstar/engine/resources/recurrent/config.py index 097b0a4f51..26ef92ec22 100644 --- a/mstar/engine/resources/recurrent/config.py +++ b/mstar/engine/resources/recurrent/config.py @@ -7,7 +7,6 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field, replace -from math import prod from typing import TYPE_CHECKING import torch @@ -24,17 +23,23 @@ class RecurrentBlockConfig: """One per-slot, per-layer tensor block. ``shape`` is opaque to the pool. ``shard_dims`` names the axes divided - across ranks. + across ranks. A ``dtype`` of None is the engine's compute dtype (what its + KV caches store), resolved by ``resolved_dtype`` and never written back. """ shape: tuple[int, ...] - dtype: torch.dtype + dtype: torch.dtype | None shard_dims: tuple[int, ...] = () def __post_init__(self): self.shape = tuple(self.shape) self._unsharded_shape = self.shape + def resolved_dtype(self, kv_dtype: torch.dtype | None) -> torch.dtype: + """The dtype the pool stores this block in; a resource planned against + the pool reads it here rather than from ``dtype``.""" + return self.dtype or kv_dtype or torch.get_default_dtype() + def shard(self, num_shards: int) -> None: from mstar.distributed.utils import divide @@ -43,14 +48,6 @@ def shard(self, num_shards: int) -> None: shape[dim] = divide(self._unsharded_shape[dim], num_shards) self.shape = tuple(shape) - @property - def numel(self) -> int: - return prod(self.shape) - - @property - def nbytes(self) -> int: - return self.numel * torch.empty((), dtype=self.dtype).element_size() - class RecurrentGeometry(ABC): @abstractmethod @@ -91,7 +88,7 @@ def conv_dim(self) -> int: def to_blocks( self, - state_dtype: torch.dtype = torch.float32, + state_dtype: torch.dtype | None = torch.float32, conv_dtype: torch.dtype = torch.bfloat16, ) -> dict[str, RecurrentBlockConfig]: """Pool blocks for this geometry; head counts are pre-sharding.""" @@ -183,7 +180,7 @@ def conv_dim(self) -> int: def to_blocks( self, - state_dtype: torch.dtype = torch.float32, + state_dtype: torch.dtype | None = torch.float32, conv_dtype: torch.dtype = torch.bfloat16, ) -> dict[str, RecurrentBlockConfig]: """Pool blocks for this geometry; head counts pre-sharding (heads and @@ -276,14 +273,6 @@ def shard(self, num_shards: int) -> None: for block in self.blocks.values(): block.shard(num_shards) - @property - def slot_bytes(self) -> int: - return self.num_layers * sum(b.nbytes for b in self.blocks.values()) - - @property - def total_bytes(self) -> int: - return self.slot_bytes * self.max_slots - @dataclass class RecurrentStateSpec(NodeResourceSpec): diff --git a/mstar/engine/resources/recurrent/pool.py b/mstar/engine/resources/recurrent/pool.py index fa34cd5beb..6194c8c00c 100644 --- a/mstar/engine/resources/recurrent/pool.py +++ b/mstar/engine/resources/recurrent/pool.py @@ -76,9 +76,14 @@ def build( config = spec.config if info.joint_comm_group is not None: config.shard(info.joint_comm_group.world_size) - return cls(device=info.device, config=config, name=spec.resource_key) - - def __init__(self, device: torch.device, config: RecurrentStateConfig, name: str = "recurrent_state"): + # resolved here, not written back: the config is the model's, shared by its specs + return cls(device=info.device, config=config, name=spec.resource_key, kv_dtype=info.kv_dtype) + + def __init__( + self, device: torch.device, config: RecurrentStateConfig, name: str = "recurrent_state", + kv_dtype: torch.dtype | None = None, + ): + """``kv_dtype`` (the engine's) stands in for a block's ``dtype=None``.""" self._device = device self.config = config # the resource key, for errors @@ -90,7 +95,7 @@ def __init__(self, device: torch.device, config: RecurrentStateConfig, name: str self._blocks: dict[str, torch.Tensor] = { name: torch.zeros( (config.num_layers, config.max_slots, *block.shape), - dtype=block.dtype, + dtype=block.resolved_dtype(kv_dtype), device=device, ) for name, block in config.blocks.items() @@ -127,12 +132,14 @@ def __init__(self, device: torch.device, config: RecurrentStateConfig, name: str self._preplanned = False self._cached_plan_output: dict[str, RecurrentAddressing] | None = None + # from the blocks, not the config: a dtype=None block is sized as resolved + total_bytes = sum(t.nbytes for t in self._blocks.values()) logger.info( "recurrent state pool: %d usable slots (+%d sink) x %d layers, " "%.2f MiB (%.2f KiB/slot), blocks=%s", config.usable_slots, config.max_slots - config.usable_slots, config.num_layers, - config.total_bytes / 2**20, config.slot_bytes / 2**10, + total_bytes / 2**20, total_bytes / config.max_slots / 2**10, {n: tuple(b.shape) for n, b in config.blocks.items()}, ) diff --git a/mstar/model/nemotron_duplex/components/audio_codec.py b/mstar/model/nemotron_duplex/components/audio_codec.py index 12f94e577c..d905462216 100644 --- a/mstar/model/nemotron_duplex/components/audio_codec.py +++ b/mstar/model/nemotron_duplex/components/audio_codec.py @@ -29,6 +29,7 @@ from torch.nn import functional as F from mstar.model.nemotron_duplex.components._util import RawWeight, param +from mstar.model.nemotron_duplex.components.dft import irfft_dim1 from mstar.model.nemotron_duplex.config import CodecConfig # The encoder/decoder layer specs — (kind, weight_shape, stride, transpose) for @@ -155,14 +156,19 @@ def _spectrogram(wav: torch.Tensor, n_fft: int, hop_length: int, win_length: int def _spec_to_wav( spec: torch.Tensor, n_fft: int, hop_length: int, win_length: int, constrain_value_range: bool = True, + by_matmul: bool = False, ) -> torch.Tensor: - """Inverse STFT via windowed overlap-add (ports NeMo ``spec_to_wav``).""" + """Inverse STFT via windowed overlap-add (ports NeMo ``spec_to_wav``). + ``by_matmul`` takes the inverse DFT as a matmul, for a captured graph (see ``dft``).""" with torch.autocast(device_type=spec.device.type, enabled=False): pad = (win_length - hop_length) // 2 T = spec.size(-1) window = torch.hann_window(win_length, device=spec.device) - ifft = torch.fft.irfft(spec, n=n_fft, dim=-2, norm="backward") + if by_matmul: + ifft = irfft_dim1(spec, n_fft) + else: + ifft = torch.fft.irfft(spec, n=n_fft, dim=-2, norm="backward") window_unsqz = window.unsqueeze(-1) if constrain_value_range: @@ -179,14 +185,38 @@ def _spec_to_wav( ifft, output_size=(1, output_size), kernel_size=(1, win_length), stride=(1, hop_length), )[..., 0, 0, pad:-pad] - window_sq = window.square().expand(T, -1).transpose(0, 1) - window_envelope = F.fold( - window_sq, output_size=(1, output_size), kernel_size=(1, win_length), stride=(1, hop_length), - ).squeeze()[pad:-pad] + # The envelope is ~0 only where the window and hop leave a gap, which + # check_window_envelope refuses when the codec is built; a device-side + # guard keeps the decode free of host syncs, for CUDA-graph capture. + env = _window_envelope(window, T, hop_length, win_length) + return torch.where(env > _ENVELOPE_FLOOR, wav / env, torch.zeros_like(wav)) + + +_ENVELOPE_FLOOR = 1e-11 + + +def _window_envelope(window: torch.Tensor, T: int, hop_length: int, win_length: int) -> torch.Tensor: + """The overlap-added squared window of ``T`` frames, trimmed as ``_spec_to_wav`` + trims the signal.""" + pad = (win_length - hop_length) // 2 + output_size = (T - 1) * hop_length + win_length + return F.fold( + window.square().expand(T, -1).transpose(0, 1), output_size=(1, output_size), + kernel_size=(1, win_length), stride=(1, hop_length), + ).squeeze()[pad:-pad] + - assert (window_envelope > 1e-11).all(), "Window envelope has zero values, cannot normalize." - wav = wav / window_envelope - return wav +def check_window_envelope(window: torch.Tensor, hop_length: int) -> None: + """Refuse a window and hop whose overlap-add has a gap at any frame count. + A sample within one window of either end sees the same frames for every + count of at least ``ceil(win / hop)``, and an interior one the full periodic + overlap, so the counts up to twice that and one more cover them all.""" + win_length = window.numel() + k = -(-win_length // hop_length) + for T in range(1, 2 * k + 2): + if not (_window_envelope(window, T, hop_length, win_length) > _ENVELOPE_FLOOR).all(): + raise ValueError(f"the codec's window overlap-adds to ~0 at hop {hop_length} ({T} frames): " + "its inverse STFT cannot normalize") class AudioCodec(nn.Module): @@ -201,6 +231,8 @@ def __init__(self, config: CodecConfig | None = None): self.hop_length = cfg.hop_length self.wav_to_token_ratio = cfg.wav_to_token_ratio self.max_magnitude = cfg.max_magnitude + # on the host: the model builds the codec under torch.device("meta") + check_window_envelope(torch.hann_window(cfg.n_fft, device="cpu"), cfg.hop_length) # ---- autoencoder halves ------------------------------------------------- def ae_encode(self, x: torch.Tensor) -> torch.Tensor: @@ -212,8 +244,9 @@ def ae_encode(self, x: torch.Tensor) -> torch.Tensor: x = self.encoder.run(x) return x.transpose(-1, -2) - def ae_decode(self, x: torch.Tensor, constrain_value_range: bool = True) -> torch.Tensor: - """Continuous latents ``[B, T', latent]`` -> waveform ``[B, 1, T]``.""" + def ae_decode(self, x: torch.Tensor, constrain_value_range: bool = True, by_matmul: bool = False) -> torch.Tensor: + """Continuous latents ``[B, T', latent]`` -> waveform ``[B, 1, T]``; ``by_matmul`` + as ``_spec_to_wav`` takes it.""" x = x.transpose(-1, -2) x = self.decoder.run(x) with torch.autocast(device_type=x.device.type, enabled=False): @@ -230,6 +263,7 @@ def ae_decode(self, x: torch.Tensor, constrain_value_range: bool = True) -> torc x = _spec_to_wav( spec, self.n_fft, self.hop_length, self.n_fft, constrain_value_range=constrain_value_range, + by_matmul=by_matmul, ).unsqueeze(1) return x @@ -248,10 +282,12 @@ def encode(self, x: torch.Tensor, x_len: torch.Tensor | None = None): code_len = x_len // self.wav_to_token_ratio if x_len is not None else None return code, code_len - def decode(self, code: torch.Tensor, code_len: torch.Tensor | None = None, constrain_value_range: bool = True): - """Codes ``[B, T', num_quantizers]`` -> (waveform ``[B, 1, T]``, wav_len).""" + def decode(self, code: torch.Tensor, code_len: torch.Tensor | None = None, constrain_value_range: bool = True, + by_matmul: bool = False): + """Codes ``[B, T', num_quantizers]`` -> (waveform ``[B, 1, T]``, wav_len); ``by_matmul`` + as ``_spec_to_wav`` takes it.""" z_q = self.dequantize(code) - x_hat = self.ae_decode(z_q, constrain_value_range=constrain_value_range) + x_hat = self.ae_decode(z_q, constrain_value_range=constrain_value_range, by_matmul=by_matmul) wav_len = code_len * self.wav_to_token_ratio if code_len is not None else None return x_hat, wav_len diff --git a/mstar/model/nemotron_duplex/components/conformer.py b/mstar/model/nemotron_duplex/components/conformer.py index 920aa35980..14139ebd7c 100644 --- a/mstar/model/nemotron_duplex/components/conformer.py +++ b/mstar/model/nemotron_duplex/components/conformer.py @@ -19,14 +19,22 @@ from __future__ import annotations import math +from typing import TYPE_CHECKING, NamedTuple import torch import torch.nn.functional as F from torch import nn +from mstar.engine.resources.kv.bounded import BoundedKVConfig, SinkWindow, StreamPosition +from mstar.engine.resources.recurrent.config import RecurrentBlockConfig from mstar.model.nemotron_duplex.components._util import param +from mstar.model.nemotron_duplex.components.dft import stft_as_real, stft_basis from mstar.model.nemotron_duplex.config import ConformerSTTConfig +if TYPE_CHECKING: + from mstar.engine.resources.kv.bounded.manager import BoundedKVManager + from mstar.engine.resources.recurrent.pool import RecurrentStatePool + # Pure numerical constants (not architecture dims / not in config.json). INF_VAL = 10000.0 # attention mask fill value (NeMo RelPos MHA) LOG_GUARD = 2.0 ** -24 # log-mel zero guard (NeMo FilterbankFeatures log_zero_guard_value) @@ -327,29 +335,61 @@ def load_weights(self, weights): # --------------------------------------------------------------------------- class StreamState: - """One call's place in its audio stream, and the encoder state that carries - it across chunks. Host counters drive every slice, so a step needs no sync. - - ``signal`` is the preemphasized audio, laid out as ``torch.stft(center=True)`` - pads it: ``n_fft // 2`` reflected samples in front, then the call. ``mel`` - holds every mel frame so far; ``k``/``v`` the last ``att_context`` keys and - values of each layer's attention, ``conv`` the last ``kernel - 1`` inputs of - each layer's depthwise conv. - """ + """One call's place in its audio stream: host counters, which drive every + slice so a step needs no sync. The audio and mel frames a later step still + reads live in the ``front`` pool (``front_state_blocks``).""" + + __slots__ = ("samples", "mels", "position") - def __init__(self, config: ConformerSTTConfig, max_samples: int, device, dtype): - c = config - self.pad = c.n_fft // 2 - self.signal = torch.zeros(max_samples + 2 * self.pad, dtype=torch.float32, device=device) - self.mel = torch.zeros(max_samples // c.hop_length + 1, c.n_mels, dtype=torch.float32, device=device) - ctx = c.att_context_size[0] - self.k = torch.zeros(c.num_layers, c.num_heads, ctx, c.head_dim, dtype=dtype, device=device) - self.v = torch.zeros_like(self.k) - self.conv = torch.zeros(c.num_layers, c.d_model, c.conv_kernel_size - 1, dtype=dtype, device=device) - self.last_sample = torch.zeros(1, dtype=torch.float32, device=device) + def __init__(self): self.samples = 0 # samples of the call so far self.mels = 0 # mel frames computed - self.frames = 0 # encoder frames emitted (= frames in the attention / conv state) + # its place in the bounded KV, which a step's declarations share + self.position = StreamPosition(source_len=0, written=0) + + @property + def frames(self) -> int: + """Encoder frames emitted.""" + return self.position.written + + @frames.setter + def frames(self, value: int): + self.position = StreamPosition(source_len=0, written=value) + + +def front_state_blocks(config: ConformerSTTConfig) -> dict[str, RecurrentBlockConfig]: + """A call's front-end state, all a later step reads: its newest + preemphasized samples (``sig``; every sample until its first mel frame, + then the window of the next one), its newest mel frames (``mel``; the + subsampling window of the next encoder frame) and its last raw sample, for + the preemphasis. Right-aligned: the newest sample / frame is last.""" + return { + "sig": RecurrentBlockConfig(shape=(config.hop_length * config.subsampling_factor,), dtype=torch.float32), + "mel": RecurrentBlockConfig(shape=(2 * config.subsampling_factor, config.n_mels), dtype=torch.float32), + "last": RecurrentBlockConfig(shape=(1,), dtype=torch.float32), + } + + +def encoder_state_blocks(config: ConformerSTTConfig, + dtype: torch.dtype | None = None) -> dict[str, RecurrentBlockConfig]: + """A call's conv state, per layer: the last ``kernel - 1`` inputs of its + depthwise conv (its attention window is the bounded KV, ``ENCODER_KV``). + By default in the engine's compute dtype, which the layers compute it in: + stored in another, a call's frames drift from the offline ones.""" + return { + "conv": RecurrentBlockConfig(shape=(config.d_model, config.conv_kernel_size - 1), dtype=dtype), + } + + +def encoder_kv_config(config: ConformerSTTConfig, max_slots: int) -> BoundedKVConfig: + """A call's attention keys and values, per layer: each new frame attends the + ``att_context_size[0]`` frames before it and itself, so the KV keeps that + many. In the engine's compute dtype, as ``encoder_state_blocks``.""" + return BoundedKVConfig( + num_layers=config.num_layers, num_heads=config.num_heads, head_dim=config.head_dim, rows_per_request=1, + max_source_len=0, retention=SinkWindow(0, config.att_context_size[0]), + query_window=config.att_context_size[0], max_slots=max_slots, dtype=None, + ) def final_mel_frames(config: ConformerSTTConfig, samples: int, final: bool) -> int: @@ -371,20 +411,43 @@ def final_encoder_frames(config: ConformerSTTConfig, mels: int, samples: int, fi return (mels - 1) // config.subsampling_factor + 1 +def device_counts(values, device: torch.device) -> torch.Tensor: + """Host counts on the device without a blocking copy, which would wait for + every kernel queued ahead of this step.""" + t = torch.tensor(values, dtype=torch.long) + return (t.pin_memory() if device.type == "cuda" else t).to(device, non_blocking=True) + + +class EncoderInputs(NamedTuple): + """A step's new encoder-input frames, ``x (B, t, D)`` zero-padded to the + most any call got, and per call its new ones.""" + x: torch.Tensor + n: list[int] + + class ConformerStreamer: """Runs a ``Perception`` incrementally over a batch of calls, matching what it computes over each whole call up to float rounding: about 1e-6 in fp32, about 1% relative in bf16, where each chunking takes different kernel shapes. - - Mel: frames are computed once their window is complete, from the padded + - Mel: frames are computed once their window is complete, from the preemphasized signal, with the featurizer's own ops. - Subsampling: the strided convs are recomputed over a window that starts 16 mel frames before the first new output frame. An output frame reads 15 mel frames, so the first two outputs of the window are discarded and the rest are exact; the window starts on a multiple of 8 so the strides line up with the offline ones. - - Layers: attention reads the cached keys/values of the last 70 frames, - the depthwise conv its cached last 8 inputs. + - Layers: each new frame attends the 70 frames before it and itself in the + bounded KV, which keeps a call's last 70 keys/values per layer; the + depthwise conv reads its cached last 8 inputs. + + A call's state is its host counters (``StreamState``), two pool slots, the + front end's (``front_state_blocks``) and the conv tails + (``encoder_state_blocks``), and its bounded-KV slot. The resources' plans + for a step come from the positions its declaration gives + (``StreamState.frames``). ``steady_step`` is a whole step for calls that + append exactly one frame of audio, all device work over fixed shapes, so it + can replay as a CUDA graph; ``front`` + ``layers`` take any step. """ def __init__(self, perception: "Perception"): @@ -392,114 +455,204 @@ def __init__(self, perception: "Perception"): self.c = perception.config enc = perception.encoder self._ctx = self.c.att_context_size[0] + self.frame_samples = self.c.hop_length * self.c.subsampling_factor + # steady_step's slices take a featurizer whose half window spans more + # than one hop and at most two (n_fft 512, hop 160); any other runs + # every step through front + layers + self.has_steady_step = -(-(self.c.n_fft // 2) // self.c.hop_length) == 2 # sinusoids of relative positions 0..ctx, row d = position d: all a query can see device = enc.layers[0].self_attn.linear_pos.weight.device pe = ConformerEncoder._rel_pos_emb(self._ctx + 1, self.c.d_model, device, torch.float32) self._pe = pe[0, : self._ctx + 1].flip(0) self._pos: dict = {} # each layer's linear_pos(_pe), by layer, dtype and autocast: fixed weights + feat = perception.preprocessor.featurizer + self._stft_basis = stft_basis(feat.window.detach(), self.c.n_fft) # steady_step's STFT, as a matmul # -- front end ----------------------------------------------------------- - def append(self, st: StreamState, chunk: torch.Tensor) -> None: - n = chunk.numel() - x = chunk.to(torch.float32).reshape(-1) - # last_sample starts at 0: the featurizer's preemphasis keeps the first sample - y = x - self.c.preemph * torch.cat([st.last_sample, x[:-1]]) - st.signal[st.pad + st.samples: st.pad + st.samples + n] = y - st.last_sample = x[-1:].clone() - st.samples += n - - def _mel_frames(self, sts: list[StreamState], lo: list[int], hi: list[int], final: list[bool]) -> None: - """Fill ``st.mel[lo:hi]`` for each call, batched.""" - c, pad, hop = self.c, self.c.n_fft // 2, self.c.hop_length - rows, idx = [], [] - for i, st in enumerate(sts): - if hi[i] <= lo[i]: - continue - if lo[i] == 0: # the reflected left edge, as torch.stft pads it - st.signal[:pad] = st.signal[pad + 1: 2 * pad + 1].flip(0) - if final[i]: # and the right edge at the end of the call - end = pad + st.samples - st.signal[end: end + pad] = st.signal[end - pad - 1: end - 1].flip(0) - rows.append(st.signal[lo[i] * hop: (hi[i] - 1) * hop + c.n_fft]) - idx.append(i) - if not rows: - return - width = max(r.numel() for r in rows) - batch = torch.stack([F.pad(r, (0, width - r.numel())) for r in rows]) - feat = self.p.preprocessor.featurizer - spec = torch.stft(batch, n_fft=c.n_fft, hop_length=hop, win_length=c.win_length, center=False, - window=feat.window.to(torch.float32), return_complex=True) - spec = torch.view_as_real(spec) + def _mel(self, batch: torch.Tensor, by_matmul: bool = False) -> torch.Tensor: + """``(B, samples)`` preemphasized windows -> ``(B, T, n_mels)`` log-mel frames. + ``by_matmul`` takes the STFT as a matmul, for a captured graph (see ``dft``); + otherwise it is the featurizer's ``torch.stft``.""" + c, feat = self.c, self.p.preprocessor.featurizer + if by_matmul: + spec = stft_as_real(batch, self._stft_basis, c.n_fft, c.hop_length) + else: + spec = torch.view_as_real(torch.stft( + batch, n_fft=c.n_fft, hop_length=c.hop_length, win_length=c.win_length, center=False, + window=feat.window.to(torch.float32), return_complex=True)) mag = torch.sqrt(spec.pow(2).sum(-1)) if c.mag_power != 1.0: mag = mag.pow(c.mag_power) mel = torch.log(torch.matmul(feat.fb.to(mag.dtype), mag) + LOG_GUARD) # (B, n_mels, T) - for row, i in enumerate(idx): - sts[i].mel[lo[i]:hi[i]] = mel[row, :, : hi[i] - lo[i]].transpose(0, 1).to(torch.float32) - - def _subsample(self, sts: list[StreamState], first: list[int], last: list[int], - mels: list[int]) -> list[torch.Tensor]: - """Encoder-input frames ``first..last`` (inclusive) of each call, from a - window of its mel frames.""" - f = self.c.subsampling_factor - starts = [max(0, f * j - 2 * f) for j in first] - windows = [sts[i].mel[starts[i]: mels[i]] for i in range(len(sts))] - lens = [w.shape[0] for w in windows] - width = max(lens) - out = torch.stack([F.pad(w, (0, 0, 0, width - w.shape[0])) for w in windows]).unsqueeze(1) # (B, 1, T, n_mels) - # A strided stage reads one frame past a row's end, where a row on its own - # reads the conv's zero padding: zero each stage's input past each row. - n = torch.tensor(lens, pin_memory=out.is_cuda).to(out.device, non_blocking=True).view(-1, 1, 1, 1) - for m in self.p.encoder.pre_encode.conv: - if isinstance(m, _CausalConv2d): - t = torch.arange(out.shape[2], device=out.device).view(1, 1, -1, 1) - out = out * (t < n).to(out.dtype) - n = n // 2 + 1 - out = m(out) # -> (B, C, T', F') + return mel.transpose(1, 2).to(torch.float32) + + def _subsample(self, windows: torch.Tensor, lens: list[int] | None = None) -> torch.Tensor: + """``(B, T, n_mels)`` mel windows -> ``(B, T', D)`` encoder-input frames; + ``lens`` gives each row's length when the rows are zero-padded to ``T``.""" + out = windows.unsqueeze(1) # (B, 1, T, n_mels) + if lens is None: + out = self.p.encoder.pre_encode.conv(out) # (B, C, T', F') + else: + # A strided stage reads one frame past a row's end, where a row on its own + # reads the conv's zero padding: zero each stage's input past each row. + n = device_counts(lens, out.device).view(-1, 1, 1, 1) + for m in self.p.encoder.pre_encode.conv: + if isinstance(m, _CausalConv2d): + t = torch.arange(out.shape[2], device=out.device).view(1, 1, -1, 1) + out = out * (t < n).to(out.dtype) + n = n // 2 + 1 + out = m(out) b, ch, t, fr = out.shape - out = self.p.encoder.pre_encode.out(out.transpose(1, 2).reshape(b, t, -1)) # (B, T', D) - return [out[i, first[i] - starts[i] // f: last[i] - starts[i] // f + 1] for i in range(len(sts))] + return self.p.encoder.pre_encode.out(out.transpose(1, 2).reshape(b, t, -1)) + + def is_steady(self, st: StreamState, chunk: torch.Tensor | None, final: bool) -> bool: + """Whether ``steady_step`` serves this call's step: one frame of audio + onto a call already two frames in, on the frame grid.""" + fs = self.frame_samples + return (self.has_steady_step and not final and chunk is not None and chunk.numel() == fs + and st.samples >= 2 * fs and st.samples % fs == 0) + + def steady_step(self, chunk: torch.Tensor, front: "RecurrentStatePool", + kv: "BoundedKVManager", pool: "RecurrentStatePool") -> torch.Tensor: + """One frame for each call, from ``chunk (B, frame_samples)``: the new + mel frames, the subsampling window and the layers, all on the device. + It is captured, so its STFT is a matmul: a graph must not hold a cuFFT + plan, which the plan cache can destroy (see ``dft``). + + A call ``S = (k - 1) fs`` samples in has ``8(k - 1) - 1`` mel frames and + ``k - 1`` encoder frames. This frame of audio adds mel frames + ``8k - 9 .. 8k - 2`` and encoder frame ``k - 1``, whose subsampling + window is the 15 mel frames before them plus them (the third output).""" + c, fs, f = self.c, self.frame_samples, self.c.subsampling_factor + last, sig = front.gather("last"), front.gather("sig") + mel_tail = front.gather("mel") + x = chunk.to(torch.float32) + y = x - c.preemph * torch.cat([last, x[:, :-1]], dim=1) + win = torch.cat([sig, y], dim=1) # samples S - fs .. S + fs + hop, pad = c.hop_length, c.n_fft // 2 + # mel frames f(k-1)-1 .. fk-2 span samples [S - hop - pad, S + fs - 2 hop + pad) + mel_new = self._mel(win[:, fs - hop - pad: 2 * fs - 2 * hop + pad], by_matmul=True) + window = torch.cat([mel_tail[:, 1:], mel_new], dim=1) # (B, 23, n_mels) + new = self._subsample(window)[:, 2:3] + front.scatter_("sig", y) + front.scatter_("last", x[:, -1:]) + front.scatter_("mel", torch.cat([mel_tail, mel_new], dim=1)[:, -2 * f:]) + return self.layers(new, kv, pool, chunk.new_ones(chunk.shape[0], dtype=torch.long)) + + def front(self, sts: list[StreamState], chunks: list[torch.Tensor | None], final: list[bool], n: list[int], + front: "RecurrentStatePool") -> EncoderInputs | None: + """Append each call's chunk; the ``n`` encoder-input frames it made final + (``new_frames``, counted before the step), or None when no call has one. + A call with no chunk and not at its end is left as it is.""" + c, pad, hop, f = self.c, self.c.n_fft // 2, self.c.hop_length, self.c.subsampling_factor + keep_sig, keep_mel = self.frame_samples, 2 * f + busy = [i for i, (ch, fin) in enumerate(zip(chunks, final, strict=True)) + if fin or (ch is not None and ch.numel())] + if not busy: + return None + last, sig = front.gather("last"), front.gather("sig") + mel_tail = front.gather("mel") + new_sig, new_last, new_mel = list(sig), list(last), list(mel_tail) + rows, lo_hi, mel_rows = [], [], [] + for i in busy: + st, chunk = sts[i], chunks[i] + held = min(st.samples, keep_sig) + win, base = sig[i, keep_sig - held:], st.samples - held # win[0] is sample ``base`` + if chunk is not None and chunk.numel(): + xr = chunk.to(torch.float32).reshape(-1) + # ``last`` reads 0 in a fresh slot: the featurizer's preemphasis keeps the first sample + y = xr - c.preemph * torch.cat([last[i], xr[:-1]]) + win = torch.cat([win, y]) + new_last[i] = xr[-1:] + st.samples += xr.numel() + new_sig[i] = F.pad(win[-keep_sig:], (max(0, keep_sig - win.numel()), 0)) + # a call's first mel frame reflects n_fft/2 samples past its first, as + # the featurizer's padding does, so it needs more than that to start + if st.samples <= pad: + continue + lo, hi = st.mels, final_mel_frames(c, st.samples, final[i]) + if hi <= lo: + continue + if lo * hop < pad: # frame lo reads the reflected left edge, as torch.stft pads it + win, base = torch.cat([win[1: pad + 1].flip(0), win]), base - pad # (base is 0: all held) + if final[i]: # and the right edge at the end of the call + win = torch.cat([win, win[-pad - 1: -1].flip(0)]) + rows.append(win[lo * hop - pad - base: (hi - 1) * hop + pad - base]) + lo_hi.append((i, lo, hi)) + added = {} + if rows: + width = max(r.numel() for r in rows) + mel = self._mel(torch.stack([F.pad(r, (0, width - r.numel())) for r in rows])) + added = {i: mel[row, : hi - lo] for row, (i, lo, hi) in enumerate(lo_hi)} + for i in busy: + st = sts[i] + if st.samples <= pad: # no mel frame yet, so no encoder frame + continue + held = min(st.mels, keep_mel) + mels, first = mel_tail[i, keep_mel - held:], st.mels - held # mels[0] is mel frame ``first`` + if i in added: + mels = torch.cat([mels, added[i]]) + st.mels += added[i].shape[0] + new_mel[i] = F.pad(mels[-keep_mel:], (0, 0, max(0, keep_mel - mels.shape[0]), 0)) + if n[i]: + start = max(0, f * st.frames - 2 * f) # the subsampling window's first mel frame + mel_rows.append((i, mels[start - first:], (st.frames - start // f, n[i]))) + front.scatter_("sig", torch.stack(new_sig)) + front.scatter_("last", torch.stack(new_last)) + front.scatter_("mel", torch.stack(new_mel)) + if not mel_rows: + return None + lens = [w.shape[0] for _, w, _ in mel_rows] + out = self._subsample(torch.stack([F.pad(w, (0, 0, 0, max(lens) - w.shape[0])) for _, w, _ in mel_rows]), + lens) + x = out.new_zeros(len(sts), max(n), out.shape[-1]) + for row, (i, _, (at, k)) in enumerate(mel_rows): + x[i, :k] = out[row, at: at + k] + return EncoderInputs(x=x, n=n) + + def new_frames(self, st: StreamState, samples: int, final: bool) -> int: + """The encoder frames a step makes final for a call at ``st`` that then + has ``samples`` of audio, counted on the host before the step runs, so + its declaration can name them and ``front`` make them. Also without a + new mel frame: the end of the call can make its last frame final.""" + c = self.c + if samples <= c.n_fft // 2: + return 0 + mels = max(st.mels, final_mel_frames(c, samples, final)) + return max(0, final_encoder_frames(c, mels, samples, final) - st.frames) + + @staticmethod + def advance(sts: list[StreamState], enc: EncoderInputs) -> None: + """Count the step's frames into each call's state, once its layers ran.""" + for st, n in zip(sts, enc.n, strict=True): + if n: + st.frames += n # -- layers -------------------------------------------------------------- - def _window(self, sts: list[StreamState], n: list[int], t: int, device) -> tuple[torch.Tensor, torch.Tensor]: - """Which keys each new frame may attend to, and their relative positions. - Keys are the ``ctx`` cached frames (slot s holds frame ``frames - ctx + s``) - then the ``t`` new ones; the same for every layer of a step.""" - ctx = self._ctx - qi = torch.arange(t, device=device).view(t, 1) - kj = torch.arange(ctx + t, device=device).view(1, ctx + t) - dist = qi + ctx - kj # (t, ctx + t) - host = torch.tensor([[st.frames for st in sts], n], dtype=torch.long) - frames, valid = host.to(device, non_blocking=True).view(2, -1, 1, 1) - outside = ((dist < 0) | (dist > ctx)).view(1, t, ctx + t) - mask = outside | (frames - ctx + kj.view(1, 1, -1) < 0) | (kj.view(1, 1, -1) - ctx >= valid) - return mask.unsqueeze(1), dist.clamp(0, ctx) - - def _attend(self, layer: "ConformerLayer", x: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, - mask: torch.Tensor, dist: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + def _attend(self, layer: "ConformerLayer", li: int, x: torch.Tensor, kv: "BoundedKVManager") -> torch.Tensor: + """Self-attention of the new frames ``x (B, t, D)``: each attends the + ``ctx`` frames before it and itself, the older ones from ``kv``, which + then keeps this step's keys and values. Transformer-XL's position term + depends on the distance alone, so it goes in as the band of distances + ``0..ctx`` (``bd_all``), unscaled: the kernel scales the sum.""" attn = layer.self_attn b, t, _ = x.shape h, dk, ctx = attn.n_heads, attn.head_dim, self._ctx q = attn.linear_q(x).view(b, t, h, dk) - k_new = attn.linear_k(x).view(b, t, h, dk).transpose(1, 2) - v_new = attn.linear_v(x).view(b, t, h, dk).transpose(1, 2) - k = torch.cat([k_cache.to(k_new.dtype), k_new], dim=2) # (B, h, ctx+t, dk) - v = torch.cat([v_cache.to(v_new.dtype), v_new], dim=2) + k = attn.linear_k(x).view(b, t, h, dk).transpose(1, 2) + v = attn.linear_v(x).view(b, t, h, dk).transpose(1, 2) key = (attn, x.dtype, torch.is_autocast_enabled(x.device.type)) p = self._pos.get(key) if p is None: # (h, ctx+1, dk) p = self._pos[key] = attn.linear_pos(self._pe.to(x.dtype)).view(ctx + 1, h, dk).transpose(0, 1) - q_u = (q + attn.pos_bias_u).transpose(1, 2) + # in k's dtype: under autocast the bias add promotes q, and the kernel takes one dtype + q_u = (q + attn.pos_bias_u).to(k.dtype).transpose(1, 2) q_v = (q + attn.pos_bias_v).transpose(1, 2) - ac = torch.matmul(q_u, k.transpose(-2, -1)) # (B, h, t, ctx+t) bd_all = torch.matmul(q_v, p.transpose(-2, -1)) # (B, h, t, ctx+1) - bd = torch.gather(bd_all, 3, dist.view(1, 1, t, ctx + t).expand(b, h, t, ctx + t)) - scores = ((ac + bd) / attn.s_d_k).masked_fill(mask, -INF_VAL) - w = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0) - out = torch.matmul(w, v).transpose(1, 2).reshape(b, t, h * dk) - return attn.linear_out(out), k, v + out = kv.attend(li, q_u, k, v, rel_bias=bd_all, ieee=k.dtype == torch.float32) + return attn.linear_out(out.reshape(b, t, h * dk)) def _conv(self, layer: "ConformerLayer", x: torch.Tensor, conv_cache: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: @@ -513,67 +666,25 @@ def _conv(self, layer: "ConformerLayer", x: torch.Tensor, z = F.silu(z) return conv.pointwise_conv2(z).transpose(1, 2), tails - def _layers(self, x: torch.Tensor, sts: list[StreamState], n: list[int]) -> torch.Tensor: - mask, dist = self._window(sts, n, x.shape[1], x.device) - k_all = torch.stack([st.k for st in sts]) # (B, L, h, ctx, dk) - v_all = torch.stack([st.v for st in sts]) - c_all = torch.stack([st.conv for st in sts]) # (B, L, D, kernel-1) - ks, vs, cs = [], [], [] + def layers(self, x: torch.Tensor, kv: "BoundedKVManager", pool: "RecurrentStatePool", + n: torch.Tensor) -> torch.Tensor: + """The layer stack over ``x (B, t, D)`` and the projection. Attention + runs on ``kv`` under the plan its step declared; the conv tails are read + and written in ``pool`` at its step's rows. ``n`` is each call's new frames, + a device ``(B,)`` count (``n <= t``; a call with none keeps its conv + tails as they were).""" + b, t, _ = x.shape + keep = self.c.conv_kernel_size - 1 + # each call keeps its newest kernel-1 conv inputs: positions n .. n + keep of the concatenation + conv_idx = (n.view(b, 1, 1) + torch.arange(keep, device=x.device).view(1, 1, keep)).expand( + b, self.c.d_model, keep) for li, layer in enumerate(self.p.encoder.layers): + conv_cache = pool.gather("conv", li) r = x + layer.FC_FACTOR * layer.feed_forward1(layer.norm_feed_forward1(x)) - a, k, v = self._attend(layer, layer.norm_self_att(r), k_all[:, li], v_all[:, li], mask, dist) - r = r + a - cv, tails = self._conv(layer, layer.norm_conv(r), c_all[:, li]) + r = r + self._attend(layer, li, layer.norm_self_att(r), kv) + cv, tails = self._conv(layer, layer.norm_conv(r), conv_cache) r = r + cv r = r + layer.FC_FACTOR * layer.feed_forward2(layer.norm_feed_forward2(r)) x = layer.norm_out(r) - ks.append(k) - vs.append(v) - cs.append(tails) - # each call keeps the newest ctx keys/values and kernel-1 conv inputs: - # positions n_i .. n_i + keep of each concatenation - ctx, keep = self._ctx, self.c.conv_kernel_size - 1 - k_cat, v_cat, c_cat = torch.stack(ks, 1), torch.stack(vs, 1), torch.stack(cs, 1) - for i, st in enumerate(sts): - st.k = k_cat[i, :, :, n[i]: n[i] + ctx].to(st.k.dtype) - st.v = v_cat[i, :, :, n[i]: n[i] + ctx].to(st.v.dtype) - st.conv = c_cat[i, :, :, n[i]: n[i] + keep].to(st.conv.dtype) - return x - - # -- one step ------------------------------------------------------------- - - def step(self, sts: list[StreamState], chunks: list[torch.Tensor | None], - final: list[bool]) -> list[torch.Tensor]: - """Append each call's chunk and return its newly final frames, ``(n_i, output_dim)``.""" - c = self.c - for st, chunk in zip(sts, chunks, strict=True): - if chunk is not None and chunk.numel(): - self.append(st, chunk) - # a call's first mel frame reflects n_fft/2 samples past its first, as - # the featurizer's padding does, so it needs more than that to start - live = [i for i, st in enumerate(sts) if st.samples > st.pad] - out: list[torch.Tensor] = [torch.zeros(0, c.output_dim) for _ in sts] - if not live: - return out - L = [sts[i] for i in live] - fin = [final[i] for i in live] - mel_hi = [final_mel_frames(c, st.samples, f) for st, f in zip(L, fin, strict=True)] - self._mel_frames(L, [st.mels for st in L], mel_hi, fin) - for st, hi in zip(L, mel_hi, strict=True): - st.mels = hi - enc_hi = [final_encoder_frames(c, st.mels, st.samples, f) for st, f in zip(L, fin, strict=True)] - ready = [k for k, st in enumerate(L) if enc_hi[k] > st.frames] - if not ready: - return out - R = [L[k] for k in ready] - first = [st.frames for st in R] - last = [enc_hi[k] - 1 for k in ready] - n = [b - a + 1 for a, b in zip(first, last, strict=True)] - new = self._subsample(R, first, last, [st.mels for st in R]) - width = max(n) - x = torch.stack([F.pad(e, (0, 0, 0, width - e.shape[0])) for e in new]) # (B, t, D) - y = self.p.proj(self._layers(x, R, n)) - for k, st in enumerate(R): - st.frames += n[k] - out[live[ready[k]]] = y[k, : n[k]] - return out + pool.scatter_("conv", tails.gather(2, conv_idx).to(conv_cache.dtype), li) + return self.p.proj(x) diff --git a/mstar/model/nemotron_duplex/components/dft.py b/mstar/model/nemotron_duplex/components/dft.py new file mode 100644 index 0000000000..3a7545eb3f --- /dev/null +++ b/mstar/model/nemotron_duplex/components/dft.py @@ -0,0 +1,75 @@ +"""Fixed-size real DFTs as matmuls, for the transforms inside captured CUDA graphs. + +``torch.stft`` / ``torch.fft`` on CUDA execute cuFFT plans from torch's +per-device plan cache, an LRU whose eviction destroys the plan. A graph replay +doesn't touch the LRU, so once enough newer transform geometries have gone +through the cache eagerly, a plan a graph captured is destroyed and the replay +computes garbage. A matmul against a precomputed basis has no plan to lose. + +The matmuls run in float64: TF32 (which ``mstar.engine`` enables) doesn't apply +to it and autocast doesn't cast it, so the result rounded to float32 is at least +as accurate as the float32 FFT, and at these sizes it costs microseconds. +""" +from __future__ import annotations + +import functools +import math + +import torch + + +def _phase(n: torch.Tensor, k: torch.Tensor, size: int) -> torch.Tensor: + """``2 pi n k / size`` in float64, with ``n k`` reduced mod ``size`` first.""" + return (n.view(-1, 1) * k.view(1, -1) % size).to(torch.float64) * (2 * math.pi / size) + + +def stft_basis(window: torch.Tensor, n_fft: int) -> torch.Tensor: + """``(win_length, 2 * (n_fft // 2 + 1))`` float64: a frame's ``win_length`` + samples -> its one-sided spectrum under ``window``, real parts then + imaginary, as ``torch.stft`` computes it with the window centred in ``n_fft``.""" + win = window.numel() + left = (n_fft - win) // 2 + n = torch.arange(left, left + win, device=window.device) + phase = _phase(n, torch.arange(n_fft // 2 + 1, device=window.device), n_fft) + w = window.to(torch.float64).view(-1, 1) + return torch.cat([w * torch.cos(phase), -w * torch.sin(phase)], dim=1) + + +def stft_as_real(x: torch.Tensor, basis: torch.Tensor, n_fft: int, hop_length: int) -> torch.Tensor: + """``torch.view_as_real(torch.stft(x, n_fft, hop_length, window=window, + center=False, return_complex=True))`` for ``x (B, samples)``, with ``basis`` + from ``stft_basis(window, n_fft)``: ``(B, n_fft // 2 + 1, T, 2)`` in x's dtype.""" + win = basis.shape[0] + left = (n_fft - win) // 2 + frames = (x.shape[-1] - n_fft) // hop_length + 1 + framed = x[:, left: left + (frames - 1) * hop_length + win].unfold(-1, win, hop_length) # (B, T, win) + with torch.autocast(device_type=x.device.type, enabled=False): + spec = torch.matmul(framed.to(torch.float64), basis).to(x.dtype) # (B, T, 2F) + return spec.unflatten(-1, (2, -1)).permute(0, 3, 1, 2) + + +@functools.cache +def irfft_basis(n: int, device: torch.device) -> torch.Tensor: + """``(n, 2 * (n // 2 + 1))`` float64: a one-sided spectrum, real parts then + imaginary -> ``torch.fft.irfft(spectrum, n, norm="backward")``. Like irfft it + ignores the imaginary parts of the DC bin and (n even) the Nyquist bin. + + Built on first use; a captured region's eager warmup makes that use before + the capture.""" + bins = n // 2 + 1 + k = torch.arange(bins, device=device) + phase = _phase(torch.arange(n, device=device), k, n) + scale = torch.full((bins,), 2.0 / n, dtype=torch.float64, device=device) + real_only = (k == 0) | ((k == n // 2) & (n % 2 == 0)) + scale = torch.where(real_only, scale / 2, scale) + return torch.cat([scale * torch.cos(phase), torch.where(real_only, 0.0, -scale * torch.sin(phase))], dim=1) + + +def irfft_dim1(spec: torch.Tensor, n: int) -> torch.Tensor: + """``torch.fft.irfft(spec, n, dim=1, norm="backward")`` for a complex + ``spec (B, n // 2 + 1, T)``: ``(B, n, T)`` in spec's real dtype.""" + assert spec.shape[1] == n // 2 + 1, f"{spec.shape[1]} bins for an inverse DFT of {n}" + parts = torch.view_as_real(spec) # (B, F, T, 2) + stacked = torch.cat([parts[..., 0], parts[..., 1]], dim=1) # (B, 2F, T) + with torch.autocast(device_type=spec.device.type, enabled=False): + return torch.matmul(irfft_basis(n, spec.device), stacked.to(torch.float64)).to(parts.dtype) diff --git a/mstar/model/nemotron_duplex/config.py b/mstar/model/nemotron_duplex/config.py index 3fed506bdd..41caa01065 100644 --- a/mstar/model/nemotron_duplex/config.py +++ b/mstar/model/nemotron_duplex/config.py @@ -30,6 +30,13 @@ TALKER_POS = "talker_pos" # positions / RoPE for the talker (per-layer theta at call time) TALKER_CFG_LABEL = "_cfg" # the two talker streams combined into one plan (cond rows, then uncond) AUDIO_IN = "audio_in" # the client stream carrying the user's audio into the Encoder +ENCODER_STATE = "encoder_state" # recurrent-state pool: each call's streaming-encoder conv tails +ENCODER_FRONT = "encoder_front" # recurrent-state pool: each call's streaming-encoder audio and mel tails +ENCODER_KV = "encoder_kv" # bounded KV: each call's streaming-encoder attention window, per layer + +# Every node's CUDA-graph batch sizes: a row per live session, up to the sessions +# a deployment serves by default (``NemotronDuplexModel.DEFAULT_SESSIONS``). +SESSION_BATCH_SIZES = [1, 2, 4, 8, 16, 32] def session_too_long(max_session_s: float) -> ValueError: diff --git a/mstar/model/nemotron_duplex/nemotron_duplex_model.py b/mstar/model/nemotron_duplex/nemotron_duplex_model.py index ba5440fcf7..c06e3d5b4c 100644 --- a/mstar/model/nemotron_duplex/nemotron_duplex_model.py +++ b/mstar/model/nemotron_duplex/nemotron_duplex_model.py @@ -32,9 +32,11 @@ decode loop consuming one encoder frame per step. The user's audio is a client stream into the Encoder, so a call can stream live over ``/generate/ws``; a one-shot upload is a stream of one chunk. ``get_node_resources`` declares -the nano's paged KV cache, attention, recurrent-state pool (Mamba-2 conv + SSM -slots), the Mamba-2 resource and the sampler, and the talker's paged KV cache, -attention and positions. Both decode steps are CUDA graphs. +the encoder's bounded KV (each call's attention window) and two recurrent-state +pools (its conv tails, and its audio / mel tails), the nano's paged KV cache, attention, recurrent-state pool (Mamba-2 +conv + SSM slots), the Mamba-2 resource and the sampler, and the talker's paged +KV cache, attention and positions. Both decode steps are CUDA graphs, and so are +the encoder's live step and the codec's steady-state decode (piecewise). Contract methods crib from ``mstar/model/orpheus/orpheus_model.py`` (single streaming LLM+codec) and ``mstar/model/qwen3_omni/qwen3_omni_model.py`` (omni, @@ -42,6 +44,7 @@ """ from __future__ import annotations +import copy import logging from pathlib import Path @@ -62,15 +65,25 @@ ResourceReqConfig, SamplerSpec, SamplingReqConfig, + apply_yaml_overrides, ) +from mstar.engine.resources.kv.bounded import BoundedKVSpec from mstar.engine.resources.linear_attn.config import LinearAttnConfig, LinearAttnSpec, LinearAttnVariant from mstar.engine.resources.recurrent import Mamba2Geometry, RecurrentStateConfig, RecurrentStateSpec from mstar.graph.base import GraphEdge, GraphNode, GraphSection, Loop, TensorPointerInfo from mstar.graph.special_destinations import EMIT_TO_CLIENT from mstar.model.base import ClientChunk, ForwardPassArgs, KVLabel, LabelLayout, Model, ProcessPromptOutput, Span from mstar.model.multimodal import PromptPart +from mstar.model.nemotron_duplex.components.conformer import ( + encoder_kv_config, + encoder_state_blocks, + front_state_blocks, +) from mstar.model.nemotron_duplex.config import ( AUDIO_IN, + ENCODER_FRONT, + ENCODER_KV, + ENCODER_STATE, MAMBA, MAMBA_STATE, NANO_ATTN, @@ -403,16 +416,20 @@ def _talker_tokenizer(self): # eartts_talker: a paged KV over its 28 Gemma3 layers with attention and # positions; each session holds two streams (labels ``main`` / ``uncond``) # and RoPE runs over the full head dim with a per-layer theta the layer - # passes at call time. The encoder and the codec declare nothing. + # passes at call time. conformer_encoder: a bounded KV, each call's + # attention window, and two recurrent-state pools, its conv tails and its + # front end's audio / mel tails (one slot per session each). The codec + # declares nothing. # ------------------------------------------------------------------- - # Sessions resident at once (one recurrent slot each; ~137 MB per slot in - # fp32 for Nemotron-H). A deployment tunes it under ``resources: mamba_state``. - # 32 is the most that keep up with real time on one H100: a tick takes - # ~62 ms at 32 sessions, and 40+ miss the 80 ms budget. - DEFAULT_MAMBA_SLOTS = 32 # sessions; the runner's padding rows take no slot (#258) + # Sessions resident at once (a slot each in mamba_state, ~137 MB in fp32 for + # Nemotron-H, and in encoder_state / encoder_front / encoder_kv; pages in + # nano_kv and talker_kv). A deployment raising it raises all six (see + # validate_config_yaml). 32 is the most that keep up with real time on one + # H100: a tick takes ~62 ms at 32 sessions, and 40+ miss the 80 ms budget. + DEFAULT_SESSIONS = 32 # padding rows take no slot (#258) # Longest session served (``model_kwargs: {max_session_s: ...}``). The - # talker KV pool is sized so DEFAULT_MAMBA_SLOTS sessions this long fit at + # talker KV pool is sized so DEFAULT_SESSIONS sessions this long fit at # once: a full pool stalls every session, not just the one that outgrew it. # A page is 128 positions: 28 layers x 128 x 16 heads x 128 (the padded # head dim) x K,V x bf16 = 28 MiB. 160 s = 2001 frames, so 37 warm-up + 2001 @@ -443,11 +460,58 @@ def max_session_samples(self) -> int: def max_session_frames(self) -> int: return self.config.stt.num_frames(self.max_session_samples) - def talker_kv_pages(self, page_size: int) -> int: - """Pages for DEFAULT_MAMBA_SLOTS sessions of max_session_frames, both CFG streams.""" - positions = self.TALKER_WARMUP_POSITIONS + self.max_session_frames + 1 + def talker_kv_pages(self, page_size: int, sessions: int | None = None) -> int: + """Pages for ``sessions`` (DEFAULT_SESSIONS) sessions of max_session_frames, both CFG streams.""" + return self._session_kv_pages(self.TALKER_WARMUP_POSITIONS, 2, page_size, sessions) + + def nano_kv_pages(self, page_size: int, sessions: int | None = None) -> int: + """Pages for ``sessions`` (DEFAULT_SESSIONS) sessions of the longest prompt and max_session_frames.""" + return self._session_kv_pages(self.max_prompt_tokens, 1, page_size, sessions) + + def _session_kv_pages(self, prefix: int, streams: int, page_size: int, sessions: int | None) -> int: + sessions = self.DEFAULT_SESSIONS if sessions is None else sessions + positions = prefix + self.max_session_frames + 1 # + 1 for the pool's sink page - return self.DEFAULT_MAMBA_SLOTS * 2 * -(-positions // page_size) + 1 + return sessions * streams * -(-positions // page_size) + 1 + + def validate_config_yaml(self, config: dict, config_path: str) -> None: + """Refuse a session cap that a per-session resource cannot hold. + + A session holds a slot in each recurrent pool (``mamba_state``, + ``encoder_state``, ``encoder_front``), in ``encoder_kv``, and its nano and + talker pages from its first step to its last, so a session past any of them waits on admit until + another ends, and its frames with it. So ``max_concurrent_requests`` + sessions must fit each; the conductor queues the rest. + """ + cap = config.get("max_concurrent_requests") + if cap is None: + logger.warning( + "%s sets no max_concurrent_requests: a session past the per-session " + "resources waits on admit until another ends", config_path, + ) + return + if not isinstance(cap, int) or isinstance(cap, bool) or cap < 1: + raise ValueError(f"{config_path}: max_concurrent_requests must be a positive int; got {cap!r}") + specs = self.get_node_resources() + apply_yaml_overrides(specs, config) + for spec in specs: + if isinstance(spec, RecurrentStateSpec): + held, knob = spec.config.usable_slots, "max_slots" + elif isinstance(spec, BoundedKVSpec): + held, knob = spec.config.max_slots, "max_slots" + elif spec.resource_key in (TALKER_KV, NANO_KV): + kv = spec.config + pages = self.talker_kv_pages if spec.resource_key == TALKER_KV else self.nano_kv_pages + per_session = pages(kv.page_size, sessions=1) - 1 + held, knob = (kv.max_num_pages - 1) // per_session, "max_num_pages" + else: + continue + if held < cap: + raise ValueError( + f"{config_path}: {cap} concurrent sessions (max_concurrent_requests) need as " + f"many in resources.{spec.resource_key}, which holds {held}; raise its {knob} " + "or lower max_concurrent_requests" + ) def get_max_output_tokens(self, **model_kwargs) -> int: # One text token per audio frame, so the audio fixes the length, within @@ -468,8 +532,30 @@ def get_node_resources(self) -> list[NodeResourceSpec]: head_dim=nano.head_dim, max_seq_len=nano.max_position_embeddings, num_qo_heads=nano.num_attention_heads, + # a session holds its prompt and every frame for life (no eviction) + max_num_pages=self.nano_kv_pages(PagedKVConfig.page_size), ) return [ + RecurrentStateSpec( + resource_key=ENCODER_STATE, nodes={"conformer_encoder"}, + config=RecurrentStateConfig( + num_layers=self.config.stt.num_layers, + blocks=encoder_state_blocks(self.config.stt), # in the engine's compute dtype + max_slots=self.DEFAULT_SESSIONS + 1, + ), + ), + RecurrentStateSpec( + resource_key=ENCODER_FRONT, nodes={"conformer_encoder"}, + config=RecurrentStateConfig( + num_layers=1, blocks=front_state_blocks(self.config.stt), + max_slots=self.DEFAULT_SESSIONS + 1, + ), + ), + # no sink slot: padding rows read and write nothing in it + BoundedKVSpec( + resource_key=ENCODER_KV, nodes={"conformer_encoder"}, + config=encoder_kv_config(self.config.stt, max_slots=self.DEFAULT_SESSIONS), + ), KVSpec(resource_key=NANO_KV, nodes={"nano_llm"}, config=kv_config), AttentionSpec( resource_key=NANO_ATTN, nodes={"nano_llm"}, @@ -486,8 +572,8 @@ def get_node_resources(self) -> list[NodeResourceSpec]: n_groups=nano.n_groups, conv_kernel_size=nano.conv_kernel, ).to_blocks(), - # + 1 for the pool's sink slot, so DEFAULT_MAMBA_SLOTS sessions fit - max_slots=self.DEFAULT_MAMBA_SLOTS + 1, + # + 1 for the pool's sink slot, so DEFAULT_SESSIONS sessions fit + max_slots=self.DEFAULT_SESSIONS + 1, ), ), LinearAttnSpec( @@ -1172,12 +1258,27 @@ def _json_config(self, local_dir: str) -> NemotronDuplexConfig: """ if getattr(self, "_json_config_cache", None) is None: try: - self._json_config_cache = NemotronDuplexConfig.from_pretrained(local_dir) + cfg = NemotronDuplexConfig.from_pretrained(local_dir) except (FileNotFoundError, KeyError) as e: logger.warning("Could not read config.json (%s); using config defaults.", e) - self._json_config_cache = self.config + cfg = self.config + self._refuse_resources_sized_otherwise(cfg) + self._json_config_cache = cfg return self._json_config_cache + def _refuse_resources_sized_otherwise(self, cfg: NemotronDuplexConfig) -> None: + """The resources are declared from the dataclass defaults, before the + checkpoint is read: refuse a config.json that sizes one otherwise, + rather than run its modules against pools of another shape.""" + loaded = copy.copy(self) + loaded.config = cfg + for declared, needed in zip(self.get_node_resources(), loaded.get_node_resources(), strict=True): + if declared != needed: + raise ValueError( + f"the checkpoint's config.json sizes resources.{declared.resource_key} as {needed}, " + f"not as declared ({declared})" + ) + def _build_and_load(self, module: torch.nn.Module, device: str, autocast_dtype, local_dir: str): """meta-build -> cast -> materialize -> load the module's checkpoint subtree. @@ -1220,10 +1321,7 @@ def _create_conformer_submodule( perception = self._build_and_load(perception, device, autocast_dtype, local_dir) rnnt_dec = self._build_and_load(rnnt_dec, device, autocast_dtype, local_dir) rnnt_joint = self._build_and_load(rnnt_joint, device, autocast_dtype, local_dir) - return ConformerEncoderSubmodule( - perception, rnnt_dec, rnnt_joint, config=self.config, - max_session_samples=self.max_session_samples, - ) + return ConformerEncoderSubmodule(perception, rnnt_dec, rnnt_joint, config=self.config) def _create_talker_submodule( self, device: str, autocast_dtype: torch.dtype | None = None, @@ -1261,4 +1359,5 @@ def _create_codec_submodule( with torch.device("meta"): codec = AudioCodec(cfg.codec) codec = self._build_and_load(codec, device, autocast_dtype, local_dir) - return AudioCodecDecoderSubmodule(codec=codec, config=self.config) + # the checkpoint's config, as the talker's: the codec decodes the codes it emits + return AudioCodecDecoderSubmodule(codec=codec, config=cfg) diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index 212d09612f..0c107cb5e8 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -2,7 +2,9 @@ Nodes, and the resources each declares (``NemotronDuplexModel.get_node_resources``): - conformer_encoder (none) 16 kHz speech -> per-frame LLM embeds (+ RNN-T) + conformer_encoder encoder_kv, 16 kHz speech -> per-frame LLM embeds (+ RNN-T); per + encoder_state, session, its attention window, its layers' conv state + encoder_front and its audio / mel tails nano_llm nano_kv, nano_attn, Nemotron-H hybrid Mamba-2 / attention / MLP (9B): mamba_state, mamba, paged KV for the 4 attention layers, recurrent-pool nano_sampler slots for the 27 Mamba-2 layers, the text sampler @@ -11,7 +13,9 @@ audio_codec (none) RVQ codes -> 22.05 kHz PCM, per-request left context Both autoregressive nodes keep their whole state in engine resources and capture -their decode step as a CUDA graph; the encoder and the codec run eager. +their decode step as a CUDA graph. The encoder keeps a call's state in its bounded +KV and two pools, and a live step (or its layer stack) replays a piecewise CUDA +graph; the codec replays one for its steady-state decode window. """ from __future__ import annotations @@ -28,6 +32,10 @@ from mstar.engine.accelerator_graph_config import ( AcceleratorGraphConfig, BatchedAcceleratorGraphConfig, + PiecewiseAcceleratorGraphConfig, + PiecewiseBatchedConfig, + PiecewiseCallInputs, + PiecewiseCaptureShape, ) from mstar.engine.resources import ( AttentionStep, @@ -38,16 +46,21 @@ SlotLease, SubmoduleStep, ) +from mstar.engine.resources.kv.bounded import BoundedKVStep from mstar.engine.resources.linear_attn.config import LinearAttnStep from mstar.engine.resources.recurrent import RecurrentStep -from mstar.model.nemotron_duplex.components.conformer import ConformerStreamer, StreamState +from mstar.model.nemotron_duplex.components.conformer import ConformerStreamer, StreamState, device_counts from mstar.model.nemotron_duplex.config import ( AUDIO_IN, + ENCODER_FRONT, + ENCODER_KV, + ENCODER_STATE, MAMBA, MAMBA_STATE, NANO_ATTN, NANO_KV, NANO_SAMPLER, + SESSION_BATCH_SIZES, TALKER_ATTN, TALKER_CFG_LABEL, TALKER_KV, @@ -96,9 +109,9 @@ class NemotronHLLMSubmodule(ARNodeSubmodule): disable_torch_compile = True # Decode batch sizes captured as CUDA graphs: one row per live session per - # 80 ms tick. The recurrent pool holds a slot per row (padding rows take - # one transiently), so the model's pool sizing covers the largest bucket. - DECODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32] + # 80 ms tick. The recurrent pool holds a slot per session; padding rows + # take none (they address the sink). + DECODE_CAPTURE_BATCH_SIZES = SESSION_BATCH_SIZES DECODE_KEY = "decode" # the steady-state capture's bucket key (cg_key_info) PROMPT_DONE_KEY = "prompt_consumed" # set once the request's first step (prompt + frame) has run @@ -426,31 +439,46 @@ class ConformerEncoderSubmodule(NodeSubmodule): Audio arrives as the client streams it, one chunk per step, and the live calls' chunks run in one batched streamer pass (``ConformerStreamer``, which describes the state a call carries across chunks and its numerics), except - that a chunk of more than a second (an upload, a large append) runs in a - pass of its own within the same engine step. + that a step holding a chunk of more than a second (an upload, a large + append) is not batched (``can_batch``): the engine then runs each call of + it on its own. A call's state is bounded whatever its length: its newest audio and + mel frames (the ``ENCODER_FRONT`` pool), the conv tails (``ENCODER_STATE``) + and the last 70 frames of each layer's attention keys and values + (``ENCODER_KV``, which also runs the attention); its host counters live in + its per-request state, and its frame count is its position in the KV. + + A call one frame of audio further on the frame grid (every live step of a + client appending 80 ms) runs whole as a CUDA graph (``STEP_REGION``). The + other calls take the front end eagerly, then the layer stack as a graph + (``LAYERS_REGION``) when each got at most one frame (a live step, the + call's end), else eagerly (a one-shot upload, a large append). In a step + whose other calls take the eager layers, or when the whole-step region has + no graph for its calls, those calls run eagerly too. Holds the perception stack plus the RNN-T decoder/joint (loaded; the - user-transcript channel is not emitted yet). Declares no resources: a - call's state lives in its per-request state. + user-transcript channel is not emitted yet). """ - # Variable-length chunks; eager. Revisit (compile, capture) once the - # one-frame step dominates. + # The captured regions are not compiled (compile=False); the eager front end + # takes variable-length chunks. disable_torch_compile = True + STEP_REGION = "conformer_step" + LAYERS_REGION = "conformer_layers" + CAPTURE_BATCH_SIZES = SESSION_BATCH_SIZES + STATE_KEY = "encoder_stream" # a call's StreamState # Longest chunk batched with other calls': a pass pads every call to its # longest chunk. One second is a judgement, not a measured crossover. BATCHED_CHUNK_S = 1.0 def __init__(self, perception: nn.Module, rnnt_decoder: nn.Module, rnnt_joint: nn.Module, - config: NemotronDuplexConfig, max_session_samples: int): + config: NemotronDuplexConfig): super().__init__() self.perception = perception self.rnnt_decoder = rnnt_decoder self.rnnt_joint = rnnt_joint self.config = config - self.max_session_samples = max_session_samples self._streamer = None def prepare_inputs( @@ -474,56 +502,250 @@ def prepare_inputs( f"the call ended with {1000 * samples / stt.sample_rate:g} ms of audio; " f"it needs more than {1000 * (stt.n_fft // 2) / stt.sample_rate:g} ms" ) + final = bool(is_final_stream_chunk) + at = st or StreamState() return NodeInputs( tensor_inputs={} if wav is None else {"chunk": wav}, - kwargs={"final": bool(is_final_stream_chunk)}, + kwargs={ + "final": final, + # whether the whole-step region can serve the call: decided + # here, once, for the declaration and the forward to share + "steady": st is not None and self.streamer.is_steady(st, wav, final), + # the call's position in ENCODER_KV and the frames this step adds + "position": at.position, + "n": self.streamer.new_frames(at, samples, final), + }, + input_seq_len=0 if wav is None else wav.numel(), # samples, for can_batch ) def preprocess(self, graph_walk, engine_inputs, inputs): return { "chunk": [inp.tensor_inputs.get("chunk") for inp in inputs], "final": [inp.kwargs.get("final", False) for inp in inputs], + "steady": [inp.kwargs.get("steady", False) for inp in inputs], + "n": [inp.kwargs.get("n", 0) for inp in inputs], } def can_batch(self, batch, model_inputs) -> bool: - return True + # A pass pads every call to its longest chunk, so a chunk of more than + # BATCHED_CHUNK_S keeps its step from being batched, and the engine runs + # each call's step on its own (declare, admit, plan, forward, commit). + # Such chunks are uploads and large appends, so it is a rare step. + longest = int(self.BATCHED_CHUNK_S * self.perception.config.sample_rate) + return len(model_inputs) == 1 or all(inp.input_seq_len <= longest for inp in model_inputs) + + def max_batch_size(self, graph_walk: str) -> int: + # The engine caps a batch at its whole-forward graphs only; past the + # widest region capture every call of the step would run eagerly. + return self.CAPTURE_BATCH_SIZES[-1] + + @staticmethod + def _declare(request_ids, keys, kv: BoundedKVStep) -> SubmoduleStep: + return SubmoduleStep( + segments=[Segment(request_id=rid, label="main", span=1) for rid in request_ids], + steps={**{key: RecurrentStep() for key in keys}, ENCODER_KV: kv}, + ) + + @staticmethod + def _kv_step(request_ids, spans: list[int], positions: dict) -> BoundedKVStep: + """Each call's new frames (``spans``) at its position in the stream.""" + return BoundedKVStep( + segments=tuple(Segment(request_id=rid, label="main", span=n) + for rid, n in zip(request_ids, spans, strict=True)), + positions=positions, + ) + + def _region_kv_step(self, request_ids) -> BoundedKVStep: + """A region's calls each add one frame. A capture's or a replay's padding + row has no stream state, so it gets no position: the resource plans it + as padding.""" + positions = {} + for rid in request_ids: + state = self.request_states.get(rid) + st = None if state is None else state.get(self.STATE_KEY) + if st is not None: + positions[rid] = st.position + return self._kv_step(request_ids, [1] * len(request_ids), positions) + + def declare_step(self, graph_walk, request_ids, inputs, slot_lease=None, piecewise_leases=None, + **kwargs) -> SubmoduleStep | None: + del graph_walk, slot_lease, kwargs + # A live step where every call is one frame on replays STEP_REGION + # over the whole batch, at the bucket its lease names, and the region's + # own step owns every resource. Each such call already holds its slots, + # from a first step declared below, so that step's admit reserves nothing. + if self.STEP_REGION in (piecewise_leases or {}) and all(inp.kwargs["steady"] for inp in inputs): + return None + # Otherwise this step reserves and plans every call's slots and new + # frames, and the eager path reads that plan. A region the forward still + # replays for some calls re-plans their rows into its own buffers, after + # the eager path has read this plan. It can't own those rows instead: its + # lease names the whole batch's bucket, not the sub-batch's, which may + # have no graph, and its admit can't refuse a call that has no slot yet. + spans = [inp.kwargs["n"] for inp in inputs] + positions = {rid: inp.kwargs["position"] for rid, inp in zip(request_ids, inputs, strict=True)} + return self._declare(request_ids, (ENCODER_FRONT, ENCODER_STATE), + self._kv_step(request_ids, spans, positions)) + + @property + def streamer(self) -> ConformerStreamer: + if self._streamer is None: + self._streamer = ConformerStreamer(self.perception) + return self._streamer + + def get_piecewise_accelerator_graph_configs( + self, device: torch.device, autocast_dtype: torch.dtype, tp_world_size: int = 1, + ) -> dict[str, PiecewiseAcceleratorGraphConfig]: + """A live step, one frame of audio per call, whole (``STEP_REGION``); and + the layer stack alone for any step with at most one frame per call + (``LAYERS_REGION``).""" + del tp_world_size + stt = self.perception.config + frame_samples = stt.hop_length * stt.subsampling_factor + + def step_inputs(shape: PiecewiseCaptureShape) -> dict[str, torch.Tensor]: + return {"chunk": torch.zeros(shape.bs, frame_samples, dtype=torch.float32, device=device)} + + def layers_inputs(shape: PiecewiseCaptureShape) -> dict[str, torch.Tensor]: + return { + "x": torch.zeros(shape.bs, 1, stt.d_model, dtype=autocast_dtype, device=device), + "n": torch.zeros(shape.bs, dtype=torch.long, device=device), + } + + def region(fn, inputs, keys, lease) -> PiecewiseBatchedConfig: + return PiecewiseBatchedConfig( + capture_fn=fn, make_static_inputs=inputs, + declare_step=lambda request_ids, seq_lens: self._declare( + request_ids, keys, self._region_kv_step(request_ids)), + seq_len=1, capture_batch_sizes=self.CAPTURE_BATCH_SIZES, compile=False, lease_before_step=lease, + ) + + regions = { + self.STEP_REGION: region(self._step_region, step_inputs, (ENCODER_FRONT, ENCODER_STATE), True), + self.LAYERS_REGION: region(self._layers_region, layers_inputs, (ENCODER_STATE,), False), + } + if not self.streamer.has_steady_step: + del regions[self.STEP_REGION] + return regions - def _stream_state(self, rid, device) -> StreamState: + def _step_region(self, call: PiecewiseCallInputs) -> dict[str, torch.Tensor]: + front, pool = call.resources[ENCODER_FRONT], call.resources[ENCODER_STATE] + return {"y": self.streamer.steady_step(call.static_inputs["chunk"], front, call.resources[ENCODER_KV], pool)} + + def _layers_region(self, call: PiecewiseCallInputs) -> dict[str, torch.Tensor]: + pool = call.resources[ENCODER_STATE] + s = call.static_inputs + return {"y": self.streamer.layers(s["x"], call.resources[ENCODER_KV], pool, s["n"])} + + def _stream_state(self, rid) -> StreamState: state = self.request_state(rid) st = state.get(self.STATE_KEY) if st is None: - dtype = torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else torch.float32 - st = StreamState(self.perception.config, self.max_session_samples, device, dtype) + st = StreamState() state.add(self.STATE_KEY, st) return st - def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, **kwargs): - if self._streamer is None: - self._streamer = ConformerStreamer(self.perception) + @staticmethod + def _to_device(chunks: list, device: torch.device) -> list: + """The step's chunks on the device in one copy: a copy per call, from + pageable memory, would each wait for its transfer.""" + host = [i for i, c in enumerate(chunks) if c is not None and c.device != device] + moved = list(chunks) + if host: + flat = torch.cat([chunks[i].reshape(-1) for i in host]) + if device.type == "cuda": + flat = flat.pin_memory() + parts = flat.to(device, non_blocking=True).split([chunks[i].numel() for i in host]) + for i, part in zip(host, parts, strict=True): + moved[i] = part + return moved + + def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, steady=None, n=None, **kwargs): rids = list(engine_inputs.request_ids) device = self.perception.proj.weight.device - states = [self._stream_state(rid, device) for rid in rids] - chunks = [c.to(device) if c is not None else None for c in chunk] - longest = int(self.BATCHED_CHUNK_S * self.perception.config.sample_rate) - alone = [c is not None and c.numel() > longest for c in chunks] - shared = [i for i, a in enumerate(alone) if not a] + streamer = self.streamer + states = [self._stream_state(rid) for rid in rids] + chunks = self._to_device(chunk, device) + runners = engine_inputs.piecewise_runners or {} frames: list[torch.Tensor | None] = [None] * len(rids) - for rows in [[i] for i, a in enumerate(alone) if a] + ([shared] if shared else []): - step = self._streamer.step([states[i] for i in rows], [chunks[i] for i in rows], [final[i] for i in rows]) - for i, f in zip(rows, step, strict=True): - frames[i] = f + + # calls one frame of audio further on the frame grid: the whole step as one graph, run last + step_runner, layers_runner = runners.get(self.STEP_REGION), runners.get(self.LAYERS_REGION) + + def layers_graph(rows) -> bool: + """Whether the layer region takes these calls' new frames: at most one each.""" + new = [i for i in rows if n[i]] + return (bool(new) and max(n[i] for i in new) == 1 and layers_runner is not None + and layers_runner.can_run(len(new))) + + if engine_inputs.step is None: + # declare_step left every resource to the region, which its lease says can replay the batch + assert all(steady), "an encoder step declared nothing but is not all steady" + steady = list(range(len(rids))) + else: + # This step's KV plan writes every call's new frames, so the region (which + # writes its calls' under its own plan) may serve the steady calls only when + # that plan goes unused: the other calls' frames, if any, all go through the + # layer region. Otherwise, or with no graph for them, they run eagerly too. + rest = [i for i, s in enumerate(steady) if not s] + steady = [i for i, s in enumerate(steady) if s] + if steady and not (step_runner is not None and step_runner.can_run(len(steady)) + and (not any(n[i] for i in rest) or layers_graph(rest))): + steady = [] + steady_chunks = [chunks[i] for i in steady] + eager_n = list(n) + for i in steady: + chunks[i], eager_n[i] = None, 0 # the eager pass leaves these calls as they are + + # The rest: the same math eagerly, reading this step's plans where it + # uses them, before any region plans its own; its layers as a graph when + # no call got more than a frame. A chunk of more than a second comes in + # a step of its own (see can_batch). + enc = None + if len(steady) < len(rids): + enc = streamer.front(states, chunks, final, eager_n, engine_inputs.resources[ENCODER_FRONT]) + if enc is not None: + new = [i for i, k in enumerate(enc.n) if k] + if layers_graph(new): + # only the calls with a frame: the region plans its own rows + y = layers_runner.run(static_inputs={"x": torch.stack([enc.x[i] for i in new]), + "n": torch.ones(len(new), dtype=torch.long, device=device)}, + request_ids=[rids[i] for i in new])["y"] + y = dict(zip(new, y, strict=True)) + else: + y = self._eager_layers(enc, engine_inputs.resources[ENCODER_KV], + engine_inputs.resources[ENCODER_STATE], device) + streamer.advance(states, enc) + for i in new: + frames[i] = y[i][: enc.n[i]] + + if steady: + y = step_runner.run(static_inputs={"chunk": torch.stack(steady_chunks)}, + request_ids=[rids[i] for i in steady])["y"] + for row, i in enumerate(steady): + st = states[i] + st.samples += streamer.frame_samples + st.mels += streamer.c.subsampling_factor + st.frames += 1 + frames[i] = y[row] # One stream item per frame: a single (T, H) tensor would be ONE item, # which the chunk policy could not slice. No check_stop, so nothing to # copy to the host for one. return BatchedModelOutput( - per_rid_outputs={rid: {"audio_frame": list(f)} for rid, f in zip(rids, frames, strict=True)}, + per_rid_outputs={rid: {"audio_frame": [] if f is None else list(f)} + for rid, f in zip(rids, frames, strict=True)}, check_stop_buffers={}, ) - def forward(self, graph_walk, engine_inputs, chunk=None, final=None, **kwargs) -> NameToTensorList: + def _eager_layers(self, enc, kv, pool, device) -> list[torch.Tensor]: + """The layer stack eagerly over ``enc``, on this step's plans: each call's new frames.""" + y = self.streamer.layers(enc.x, kv, pool, device_counts(enc.n, device)) + return [row[:n] for row, n in zip(y, enc.n, strict=True)] + + def forward(self, graph_walk, engine_inputs, **kwargs) -> NameToTensorList: # one call's step: preprocess's lists of one row rid = engine_inputs.request_ids[0] - return self.forward_batched(graph_walk, engine_inputs, chunk=chunk, final=final).per_rid_outputs[rid] + return self.forward_batched(graph_walk, engine_inputs, **kwargs).per_rid_outputs[rid] class EarTTSTalkerSubmodule(ARNodeSubmodule): @@ -549,7 +771,7 @@ class EarTTSTalkerSubmodule(ARNodeSubmodule): GEN_KEY = "talker_gen" DECODE_KEY = "decode" # the steady-state capture's bucket key (cg_key_info) SPEAKER = "Aria" - DECODE_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32] + DECODE_CAPTURE_BATCH_SIZES = SESSION_BATCH_SIZES def __init__(self, talker: nn.Module, config: NemotronDuplexConfig, subword_to_char: dict | None = None, char_pad_idx: int | None = None): @@ -773,13 +995,20 @@ class AudioCodecDecoderSubmodule(NodeSubmodule): recomputed with its context: the emitted stream is the full-history decode sample for sample; the final chunk emits its held tail too. Requests whose windows have the same length are decoded in one batched - call. Declares no resources. + call; a steady-state window (left context + one chunk, every chunk after a + call's first) replays a CUDA graph (``DECODE_REGION``). Calls that started + at different times reach their chunk boundaries on different frames, so a + live load runs many small decodes. Declares no resources. """ - # Per-request context is keyed by request id -> eager; fp32 codec. + # Per-request context is keyed by request id -> eager around the decode. No + # autocast: the codec runs in the dtype it was loaded in (the engine's). disable_torch_compile = True disable_autocast = True + DECODE_REGION = "codec_decode" + CAPTURE_BATCH_SIZES = SESSION_BATCH_SIZES + CONTEXT_KEY = "codec_ctx" TAIL_KEY = "codec_tail" # the last chunk's held-back samples, int16 @@ -820,6 +1049,39 @@ def preprocess(self, graph_walk, engine_inputs, inputs): def can_batch(self, batch, model_inputs) -> bool: return True + def max_batch_size(self, graph_walk: str) -> int: + # as the encoder's: the region's widest capture, which the engine's cap does not see + return self.CAPTURE_BATCH_SIZES[-1] + + @property + def _steady_window(self) -> int: + eartts = self.config.eartts + return eartts.codec_left_context_frames + eartts.codec_chunk_frames + + def get_piecewise_accelerator_graph_configs( + self, device: torch.device, autocast_dtype: torch.dtype, tp_world_size: int = 1, + ) -> dict[str, PiecewiseAcceleratorGraphConfig]: + del autocast_dtype, tp_world_size + window, num_q = self._steady_window, self.config.eartts.num_quantizers + + def make_static_inputs(shape: PiecewiseCaptureShape) -> dict[str, torch.Tensor]: + return {"codes": torch.zeros(shape.bs, window, num_q, dtype=torch.long, device=device)} + + return { + self.DECODE_REGION: PiecewiseBatchedConfig( + capture_fn=self._decode_region, + make_static_inputs=make_static_inputs, + seq_len=1, + capture_batch_sizes=self.CAPTURE_BATCH_SIZES, + compile=False, + ), + } + + def _decode_region(self, call: PiecewiseCallInputs) -> dict[str, torch.Tensor]: + # captured: the inverse DFT as a matmul, since the cuFFT plan cache can destroy a captured plan + audio, _ = self.codec.decode(call.static_inputs["codes"], by_matmul=True) + return {"wav": audio.squeeze(1)} + def _window(self, rid: str, codes: torch.Tensor) -> tuple[torch.Tensor, int]: """This request's ``context + new`` code window and its context length.""" prev = self.request_state(rid).get(self.CONTEXT_KEY) # (n_ctx, num_q) or None @@ -847,11 +1109,14 @@ def forward_batched( by_len: dict[int, list[int]] = {} for i, (full, _) in windows.items(): by_len.setdefault(full.shape[0], []).append(i) + runner = (engine_inputs.piecewise_runners or {}).get(self.DECODE_REGION) for tf, idxs in by_len.items(): batch = torch.stack([windows[i][0] for i in idxs]).long() # (G, Tf, num_q) - code_len = torch.full((len(idxs),), tf, device=batch.device) - audio, _ = self.codec.decode(batch, code_len) # (G, 1, samples) - wav = audio.squeeze(1) # (G, samples) + if runner is not None and tf == self._steady_window and runner.can_run(len(idxs)): + wav = runner.run(static_inputs={"codes": batch}, real_bs=len(idxs))["wav"] + else: + audio, _ = self.codec.decode(batch) # (G, 1, samples) + wav = audio.squeeze(1) # (G, samples) spf = wav.shape[1] // tf # samples per frame for row, i in enumerate(idxs): full, n_ctx = windows[i] diff --git a/test/modular/test_client_stream_api.py b/test/modular/test_client_stream_api.py index fe354dca50..f09a39da4a 100644 --- a/test/modular/test_client_stream_api.py +++ b/test/modular/test_client_stream_api.py @@ -1611,3 +1611,13 @@ def test_max_call_input_mb_reads_in_mib(): server.max_call_input_bytes = entrypoint._call_input_limit({"max_call_input_mb": 0.5}) with pytest.raises(ValueError, match=r"passed 0\.5 MiB"): server._check_call_input(1 << 20) + + +def test_websocket_deflate_is_a_serving_setting(): + """Per-message deflate on the WebSocket routes is the deployment's choice: + on unless the serving config turns it off (as the duplex configs do, where + compressing live audio cost the event loop more than it saved).""" + app = object() + assert entrypoint._uvicorn_config(app, "127.0.0.1", 8000, {}).ws_per_message_deflate is True + off = entrypoint._uvicorn_config(app, "127.0.0.1", 8000, {"ws_per_message_deflate": False}) + assert off.ws_per_message_deflate is False diff --git a/test/modular/test_mamba2_geometry.py b/test/modular/test_mamba2_geometry.py index 99612584ca..6051c70c6e 100644 --- a/test/modular/test_mamba2_geometry.py +++ b/test/modular/test_mamba2_geometry.py @@ -5,11 +5,7 @@ import pytest import torch -from mstar.engine.resources.recurrent import ( - DeltaNetGeometry, - Mamba2Geometry, - RecurrentStateConfig, -) +from mstar.engine.resources.recurrent import DeltaNetGeometry, Mamba2Geometry NEMOTRON_H = Mamba2Geometry(num_heads=128, head_dim=80, state_size=128, n_groups=8, conv_kernel_size=4) @@ -19,10 +15,6 @@ def test_nemotron_h_blocks(): assert NEMOTRON_H.d_inner == 10240 and NEMOTRON_H.conv_dim == 12288 assert blocks["ssm"].shape == (128, 80, 128) and blocks["ssm"].dtype is torch.float32 assert blocks["conv"].shape == (12288, 3) and blocks["conv"].dtype is torch.bfloat16 - # per-slot bytes: fp32 ssm (5.0 MiB/layer) + bf16 conv (72 KiB/layer) - cfg = RecurrentStateConfig(num_layers=27, blocks=blocks, max_slots=64) - assert cfg.slot_bytes == 27 * (128 * 80 * 128 * 4 + 12288 * 3 * 2) - assert round(cfg.slot_bytes / 2**20) == 137 def test_blocks_round_trip(): diff --git a/test/modular/test_mamba2_resource.py b/test/modular/test_mamba2_resource.py index 4edf5169b4..519cf75c1b 100644 --- a/test/modular/test_mamba2_resource.py +++ b/test/modular/test_mamba2_resource.py @@ -108,6 +108,19 @@ def test_build_dispatches_on_the_variant_and_reads_the_geometry(): assert pool.block("ssm", 1).shape == (8, 4, 8, 16) and pool.block("conv", 0).shape == (8, GEOM.conv_dim, 3) +def test_a_state_in_the_engine_dtype_reaches_the_manager_as_the_pool_stores_it(): + spec = RecurrentStateSpec(POOL, {"llm"}, RecurrentStateConfig( + num_layers=NUM_LAYERS, blocks=GEOM.to_blocks(state_dtype=None), max_slots=8, + )) + info = EngineResourceInfo(device=torch.device("cpu"), kv_dtype=torch.float16, dependencies={POOL: spec}) + pool = RecurrentStatePool.build(spec, info) + manager = LinearAttnManager.build( + LinearAttnSpec(MAMBA, {"llm"}, LinearAttnConfig(recurrent_state=POOL, variant=LinearAttnVariant.MAMBA2)), + info, + ) + assert manager.state_dtype is pool.block("ssm", 0).dtype is torch.float16 + + def test_plan_is_a_narrow_of_the_pools_addressing(): pool, manager = build() _, c = plan(pool, manager, ["a", "b", "c"], [1, 1, 1]) diff --git a/test/modular/test_nemotron_duplex_model.py b/test/modular/test_nemotron_duplex_model.py index 1e0cc8c6fc..e035acb029 100644 --- a/test/modular/test_nemotron_duplex_model.py +++ b/test/modular/test_nemotron_duplex_model.py @@ -23,10 +23,14 @@ SamplingReqConfig, resolve_spec_dependencies, ) +from mstar.engine.resources.kv.bounded import BoundedKVSpec from mstar.engine.resources.linear_attn.config import LinearAttnSpec, LinearAttnVariant from mstar.engine.resources.recurrent import Mamba2Geometry, RecurrentStateSpec from mstar.model.base import ClientChunk from mstar.model.nemotron_duplex.config import ( + ENCODER_FRONT, + ENCODER_KV, + ENCODER_STATE, MAMBA, MAMBA_STATE, NANO_ATTN, @@ -72,13 +76,32 @@ def test_duplex_node_resources(): planned on it (NoPE: no position resource), a recurrent pool holding the 27 Mamba-2 layers' conv + SSM state with the Mamba-2 resource planned on it, and one text sampler; the talker a paged KV over its 28 layers with attention - and RoPE positions. The encoder and the codec own no engine resource.""" + and RoPE positions; the encoder a bounded KV holding each call's attention + window, and two recurrent pools, its conv tails and its front end's audio / + mel tails, a slot per session each. The codec owns no engine resource.""" model = _make_model() specs = model.get_node_resources() by_key = resolve_spec_dependencies(specs) # unique keys, dependencies satisfied - assert set(by_key) == {NANO_KV, NANO_ATTN, MAMBA_STATE, MAMBA, NANO_SAMPLER, TALKER_KV, TALKER_ATTN, TALKER_POS} + assert set(by_key) == {ENCODER_FRONT, ENCODER_STATE, ENCODER_KV, NANO_KV, NANO_ATTN, MAMBA_STATE, MAMBA, + NANO_SAMPLER, TALKER_KV, TALKER_ATTN, TALKER_POS} nano_keys = {NANO_KV, NANO_ATTN, MAMBA_STATE, MAMBA, NANO_SAMPLER} - assert all(spec.nodes == ({"nano_llm"} if spec.resource_key in nano_keys else {"eartts_talker"}) for spec in specs) + node = {**dict.fromkeys((ENCODER_STATE, ENCODER_FRONT, ENCODER_KV), {"conformer_encoder"}), + **dict.fromkeys(nano_keys, {"nano_llm"})} + assert all(spec.nodes == node.get(spec.resource_key, {"eartts_talker"}) for spec in specs) + + stt = model.config.stt + enc = by_key[ENCODER_STATE] + assert isinstance(enc, RecurrentStateSpec) and enc.config.num_layers == stt.num_layers + assert set(enc.config.blocks) == {"conv"} + assert enc.config.blocks["conv"].shape == (stt.d_model, stt.conv_kernel_size - 1) + assert enc.config.usable_slots == model.DEFAULT_SESSIONS + kv = by_key[ENCODER_KV] + assert isinstance(kv, BoundedKVSpec) + # the checkpoint's: 24 layers of 8 heads of 128, each frame over the 70 before it + dims = (kv.config.num_layers, kv.config.num_heads, kv.config.head_dim, kv.config.query_window) + assert dims == (stt.num_layers, stt.num_heads, stt.head_dim, stt.att_context_size[0]) == (24, 8, 128, 70) + assert kv.config.window_capacity == 70 and kv.config.sink_capacity == 0 + assert kv.config.max_slots == model.DEFAULT_SESSIONS and kv.config.dtype is None eartts = model.config.eartts tkv = by_key[TALKER_KV] @@ -100,7 +123,7 @@ def test_duplex_node_resources(): dims = (geom.num_heads, geom.head_dim, geom.state_size, geom.n_groups, geom.conv_kernel_size) assert dims == (128, 80, 128, 8, 4) assert pool.config.blocks["ssm"].dtype is torch.float32 and pool.config.blocks["conv"].shape == (12288, 3) - assert pool.config.usable_slots == model.DEFAULT_MAMBA_SLOTS # the sink is extra; padding rows take none + assert pool.config.usable_slots == model.DEFAULT_SESSIONS # the sink is extra; padding rows take none mamba = by_key[MAMBA] assert isinstance(mamba, LinearAttnSpec) and mamba.config.variant is LinearAttnVariant.MAMBA2 assert mamba.config.recurrent_state == MAMBA_STATE and mamba.depends_on() == {MAMBA_STATE} @@ -119,6 +142,144 @@ def test_duplex_node_resources(): assert isinstance(sampler, SamplerSpec) and sampler.vocab_size == model.config.vocab_size +def test_duplex_shipped_configs_hold_their_session_cap(): + import yaml + + model = _make_model() + for name in ("nemotron_duplex.yaml", "nemotron_duplex_disagg.yaml"): + path = CONFIG_PATH.parent / name + model.validate_config_yaml(yaml.safe_load(path.read_text()), str(path)) + + +def test_duplex_shipped_configs_turn_websocket_deflate_off(): + """Compressing live calls' appends cost the server's event loop more than + the bandwidth was worth at 32 calls, where it was measured.""" + import yaml + + for name in ("nemotron_duplex.yaml", "nemotron_duplex_disagg.yaml"): + assert yaml.safe_load((CONFIG_PATH.parent / name).read_text())["ws_per_message_deflate"] is False + + +@pytest.mark.parametrize("key", [MAMBA_STATE, ENCODER_STATE, ENCODER_FRONT, ENCODER_KV, TALKER_KV, NANO_KV]) +def test_duplex_refuses_a_session_cap_one_per_session_resource_cannot_hold(key): + """Every session holds a slot in each recurrent pool and its nano and talker + pages from start to end; a session past any of them waits on admit until + another ends. Raising the cap means raising them all, and a deployment + that misses one is refused at startup, naming it.""" + model = _make_model() + resources = {MAMBA_STATE: {"max_slots": 49}, ENCODER_STATE: {"max_slots": 49}, + ENCODER_FRONT: {"max_slots": 49}, ENCODER_KV: {"max_slots": 48}, + TALKER_KV: {"max_num_pages": model.talker_kv_pages(128, 48)}, + NANO_KV: {"max_num_pages": model.nano_kv_pages(128, 48)}} + model.validate_config_yaml({"max_concurrent_requests": 48, "resources": resources}, "c.yaml") + del resources[key] + with pytest.raises(ValueError, match=f"resources.{key}"): + model.validate_config_yaml({"max_concurrent_requests": 48, "resources": resources}, "c.yaml") + + +@pytest.mark.parametrize("kwargs", [{"max_session_s": 640.0}, {"max_prompt_tokens": 8192}]) +def test_duplex_kv_pools_hold_the_session_cap_at_the_calls_it_serves(kwargs): + """A session holds its nano pages (prompt and every frame) and its talker + pages for life, so both pools are sized from the longest call and prompt + served; one sized below the cap is refused at startup.""" + model = _make_model() + for name, value in kwargs.items(): + setattr(model, name, value) + model.validate_config_yaml({"max_concurrent_requests": 32}, "c.yaml") + with pytest.raises(ValueError, match=f"resources.{NANO_KV}"): + model.validate_config_yaml({"max_concurrent_requests": 32, "resources": {NANO_KV: {"max_num_pages": 2}}}, + "c.yaml") + + +@pytest.mark.parametrize("autocast_dtype", [torch.bfloat16, torch.float16, None]) +def test_duplex_encoder_state_is_stored_in_the_engines_compute_dtype(autocast_dtype): + """The encoder stores each step's layer state and keys / values for the + next; in another dtype than the layers compute them in, live frames drift + from one-shot ones. So the pool and the KV take the dtype the engine + resolved (``autocast_dtype: null`` computes in fp32), and no deployment + dtype is refused for them.""" + from mstar.engine.resources.base import EngineResourceInfo + from mstar.engine.resources.kv.bounded.manager import BoundedKVManager + from mstar.engine.resources.recurrent.pool import RecurrentStatePool + + model = _make_model() + model.validate_config_yaml({"autocast_dtype": autocast_dtype}, "c.yaml") + specs = {s.resource_key: s for s in model.get_node_resources()} + info = EngineResourceInfo(device=torch.device("cpu"), kv_dtype=autocast_dtype) + pool = RecurrentStatePool.build(specs[ENCODER_STATE], info) + want = torch.float32 if autocast_dtype is None else autocast_dtype + assert {pool.block(name, 0).dtype for name in specs[ENCODER_STATE].config.blocks} == {want} + assert BoundedKVManager.build(specs[ENCODER_KV], info).layer(0).dtype == want + + +class _LoadReachedError(Exception): + """Raised where the model would load a node's weights: its modules were built.""" + + +def _stub_checkpoint(monkeypatch, config=None): + """The model's checkpoint access, stubbed: no snapshot, ``config`` as its + config.json, and each node's modules returned as built.""" + import mstar.model.nemotron_duplex.nemotron_duplex_model as ndm + + monkeypatch.setattr(ndm, "_resolve_local_hf_snapshot", lambda *a, **k: "/nonexistent") + monkeypatch.setattr(NemotronDuplexConfig, "from_pretrained", classmethod(lambda cls, d: config or cls())) + + +@pytest.mark.parametrize("node", sorted(NODES)) +def test_duplex_builds_every_node_on_meta(monkeypatch, node): + """Each node's modules are built under ``torch.device("meta")`` before their + weights load, so no constructor may read a tensor's values.""" + model = _make_model() + model.model_path_hf, model.cache_dir = "unused", None + _stub_checkpoint(monkeypatch) + + def built(self, module, device, autocast_dtype, local_dir): + assert next(module.parameters()).device.type == "meta" + raise _LoadReachedError + + monkeypatch.setattr(NemotronDuplexModel, "_build_and_load", built) + with pytest.raises(_LoadReachedError): + model.get_submodule(node, device="cpu") + + +def test_duplex_refuses_a_checkpoint_its_resources_were_not_sized_for(monkeypatch): + """The resources are declared from the config defaults, before the + checkpoint is read: a config.json that sizes one differently (another + encoder depth, here) is refused when the checkpoint is read, rather than + run against pools of another shape.""" + ckpt = NemotronDuplexConfig() + ckpt.stt.num_layers = 12 + _stub_checkpoint(monkeypatch, ckpt) + with pytest.raises(ValueError, match=f"resources.{ENCODER_STATE}"): + _make_model()._json_config("/nonexistent") + + +def test_duplex_checks_the_encoder_kv_against_the_checkpoint(monkeypatch): + """The encoder KV's declaration from the defaults equals the one from a + checkpoint of the same shape (its retention is a SinkWindow, equal by value), and + a checkpoint with another attention window is refused, naming it.""" + _stub_checkpoint(monkeypatch) + _make_model()._json_config("/nonexistent") + ckpt = NemotronDuplexConfig() + ckpt.stt.att_context_size = (64, 0) + _stub_checkpoint(monkeypatch, ckpt) + with pytest.raises(ValueError, match=f"resources.{ENCODER_KV}"): + _make_model()._json_config("/nonexistent") + + +def test_duplex_codec_sizes_its_region_from_the_checkpoint(monkeypatch): + """The codec decodes the codes the talker emits, so it takes the checkpoint's + config (codebooks per frame, here), as the talker does.""" + ckpt = NemotronDuplexConfig() + ckpt.eartts.num_quantizers = 8 + _stub_checkpoint(monkeypatch, ckpt) + model = _make_model() + model.model_path_hf, model.cache_dir = "unused", None + monkeypatch.setattr(NemotronDuplexModel, "_build_and_load", lambda self, module, *a: module) + codec = model.get_submodule("audio_codec", device="cpu") + assert codec.config.eartts.num_quantizers == 8 + + def test_duplex_request_resource_configs(): model = _make_model() cfg = model.get_request_resource_configs({}, None) @@ -374,7 +535,7 @@ def test_duplex_appended_audio_becomes_stream_items_in_order(): def test_duplex_an_upload_or_append_is_one_stream_item(): """A file streams whole, so a one-file upload is one encoder step, as the - offline encoder ran it (the encoder runs a long item in a pass of its own).""" + offline encoder ran it (a long item is not batched with other calls').""" model = _make_model() long = torch.rand(20 * model.config.stt.sample_rate + 777) - 0.5 for items in ( diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index a1165b574f..5054832e0b 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -12,9 +12,23 @@ from torch import nn from mstar.engine.resources import AttentionStep, KVStep, SamplerStep +from mstar.engine.resources.base import EngineResourceInfo +from mstar.engine.resources.kv.bounded import BoundedKVSpec +from mstar.engine.resources.kv.bounded.manager import BoundedKVManager from mstar.engine.resources.linear_attn.config import LinearAttnStep -from mstar.engine.resources.recurrent import RecurrentStep +from mstar.engine.resources.recurrent import RecurrentStateConfig, RecurrentStateSpec, RecurrentStep +from mstar.engine.resources.recurrent.pool import RecurrentStatePool +from mstar.engine.resources.runner import StepRunner +from mstar.engine.resources.step import StepContext +from mstar.model.nemotron_duplex.components.conformer import ( + encoder_kv_config, + encoder_state_blocks, + front_state_blocks, +) from mstar.model.nemotron_duplex.config import ( + ENCODER_FRONT, + ENCODER_KV, + ENCODER_STATE, MAMBA, MAMBA_STATE, NANO_ATTN, @@ -28,7 +42,7 @@ ConformerEncoderSubmodule, NemotronHLLMSubmodule, ) -from mstar.model.submodule_base import BatchedModelOutput +from mstar.model.submodule_base import BatchedModelOutput, ModelInputsFromEngine H = 8 @@ -198,7 +212,7 @@ def test_nano_captures_the_decode_step(): assert isinstance(cfg, BatchedAcceleratorGraphConfig) assert cfg.capture_graph_walk == "decode" and cfg.replay_graph_walks == ["decode"] assert cfg.compile is False and cfg.additional_key_info == nano.DECODE_KEY - assert max(cfg.capture_batch_sizes) <= NemotronDuplexModel.DEFAULT_MAMBA_SLOTS + assert max(cfg.capture_batch_sizes) <= NemotronDuplexModel.DEFAULT_SESSIONS row = cfg.single_request_inputs assert row.input_seq_len == 1 and row.kwargs["mode"] == "frame" assert row.tensor_inputs["audio_frame"].shape == (1, nano.config.nano.hidden_size) @@ -276,7 +290,7 @@ def __init__(self): super().__init__() self.calls = 0 - def decode(self, codes, code_len): + def decode(self, codes, code_len=None): self.calls += 1 b, tf = codes.shape[0], codes.shape[1] per_frame = codes[:, :, 0].float() / 1000.0 # (B, Tf) @@ -293,12 +307,12 @@ def _codes(n_frames, value=0): def _run_codec(codec, rid, n_frames, value=0): - eng = SimpleNamespace(request_ids=[rid]) + eng = SimpleNamespace(request_ids=[rid], piecewise_runners={}) return codec.forward("codec_chunk", eng, codes=_codes(n_frames, value))["audio_chunk"][0] def _run_codec_batch(codec, rids, frames): - eng = SimpleNamespace(request_ids=list(rids)) + eng = SimpleNamespace(request_ids=list(rids), piecewise_runners={}) inputs = [codec.prepare_inputs("codec_chunk", None, {"codec_tokens": [_codes(n, v)]}) for n, v in frames] pre = codec.preprocess("codec_chunk", eng, inputs) @@ -326,7 +340,7 @@ def test_codec_runs_an_empty_final_chunk_as_a_flush(): inp = codec.prepare_inputs("codec_chunk", None, chunk, is_final_stream_chunk=True) assert inp is not None and inp.kwargs["flush"] and not inp.tensor_inputs _run_codec(codec, "a", 5, 1) - eng = SimpleNamespace(request_ids=["a", "b"]) + eng = SimpleNamespace(request_ids=["a", "b"], piecewise_runners={}) inputs = [codec.prepare_inputs("codec_chunk", None, {}, is_final_stream_chunk=True), codec.prepare_inputs("codec_chunk", None, {"codec_tokens": [_codes(5, 2)]})] out = codec.forward_batched("codec_chunk", eng, **codec.preprocess("codec_chunk", eng, inputs)) @@ -382,12 +396,58 @@ def test_codec_cleanup_clears_per_request_context(): assert "r" not in codec.request_states +def _captured_on_cpu(config, label, resources=None, step_runner=None, bs=(1, 2, 4)): + """The engine's piecewise runner for ``config``, captured on CPU: the real + runner (bucket padding, its own static buffers per bucket, the region's + admit / plan / commit), with a "graph" whose replay re-runs the region + eagerly. ``calls`` counts its replays.""" + from unittest import mock + + from mstar.engine import accelerator_graph_runner as cgr + from mstar.engine.accelerator_graph_backend import AcceleratorGraphPool + + class _Replay: + def __init__(self, fn): + self.fn = fn + + def replay(self): + self.fn() + + config.capture_batch_sizes = list(bs) + backend = SimpleNamespace(is_available=lambda: True, set_device=lambda: None, synchronize=lambda: None) + runner = cgr.PiecewiseAcceleratorGraphRunner( + label=label, config=config, resources=resources or {}, step_runner=step_runner or StepRunner({}), + device=torch.device("cpu"), autocast_dtype=None, num_slots=1, node_name=label, + memory_pool=AcceleratorGraphPool("cpu"), device_capability=SimpleNamespace(graph_backend=backend), + ) + # a CPU device has no capture plan, so every bucket is warmed up and captured; the warm-up's + # memory readings are the accelerator's, which a CPU device has none of + no_memory_stats = dict(memory_allocated=lambda device=None: 0, max_memory_allocated=lambda device=None: 0, + reset_peak_memory_stats=lambda device=None: None, empty_cache=lambda: None) + with mock.patch.object(cgr, "capture_into_graph", lambda fn, *a: (_Replay(fn), fn())), \ + mock.patch.multiple(torch.accelerator, **no_memory_stats): + runner.prepare_for_capture() + runner.warmup_and_capture(None) + assert runner.any_graphs and not runner.dropped_shapes + runner.calls = 0 + run = runner.run + + def counted(*args, **kwargs): + runner.calls += 1 + return run(*args, **kwargs) + + runner.run = counted + return runner + + +@pytest.mark.parametrize("region", [False, True]) @pytest.mark.parametrize("frames", [23, 25]) -def test_codec_stream_is_the_full_history_decode_with_the_real_decoder(frames): +def test_codec_stream_is_the_full_history_decode_with_the_real_decoder(frames, region): """With the real (tiny-channel, real kernel / rates / block count) decoder, the chunked stream -- 3 frames of left context, right edge carried into the next chunk -- equals the one-shot decode of all frames, sample for sample. - The final chunk flushes the held tail; 25 frames end on an empty one.""" + The final chunk flushes the held tail; 25 frames end on an empty one. With + the decode region, the steady-state windows go through it, padded.""" from mstar.model.nemotron_duplex.components.audio_codec import AudioCodec from mstar.model.nemotron_duplex.config import CodecConfig @@ -400,6 +460,12 @@ def test_codec_stream_is_the_full_history_decode_with_the_real_decoder(frames): cfg = NemotronDuplexConfig() codec = AudioCodecDecoderSubmodule(codec=real, config=cfg) assert cfg.eartts.codec_left_context_frames == 3 and codec._hold == ccfg.n_fft - ccfg.hop_length == 12 + runners = {} + if region: + cfg.eartts.num_quantizers = ccfg.num_quantizers + label = AudioCodecDecoderSubmodule.DECODE_REGION + runners[label] = _captured_on_cpu( + codec.get_piecewise_accelerator_graph_configs(torch.device("cpu"), torch.float32)[label], label) chunk = 5 codes = torch.randint(0, ccfg.codebook_size, (frames, ccfg.num_quantizers)) with torch.no_grad(): @@ -411,14 +477,88 @@ def test_codec_stream_is_the_full_history_decode_with_the_real_decoder(frames): piece = codes[start:start + chunk] inp = codec.prepare_inputs("codec_chunk", None, {"codec_tokens": [piece] if len(piece) else []}, is_final_stream_chunk=start == starts[-1]) - eng = SimpleNamespace(request_ids=["r"], resources={}, per_request_states=None) + eng = SimpleNamespace(request_ids=["r"], resources={}, per_request_states=None, + piecewise_runners=runners) out = codec.forward_batched("codec_chunk", eng, **codec.preprocess("codec_chunk", eng, [inp])) pieces.append(out["r"]["audio_chunk"][0]) stream = torch.cat(pieces) + if region: + assert runners[AudioCodecDecoderSubmodule.DECODE_REGION].calls >= 3 assert stream.shape[0] == whole.shape[0] torch.testing.assert_close(stream, whole, atol=1, rtol=0) # int16: one quantization step +def test_codec_region_replays_three_windows_with_a_padding_row(): + """Three calls in step decode their steady windows together through the + region, in a bucket of four: the padding row decodes nothing any call + receives, and each call's stream is its full-history decode.""" + from mstar.model.nemotron_duplex.components.audio_codec import AudioCodec + from mstar.model.nemotron_duplex.config import CodecConfig + + torch.manual_seed(0) + ccfg = CodecConfig(base_hidden_size=8, latent_size=16, codebook_size=16, num_quantizers=4) + real = AudioCodec(ccfg).eval() + for prm in real.parameters(): + if prm.is_floating_point(): + prm.data.normal_(0, 0.05) + cfg = NemotronDuplexConfig() + cfg.eartts.num_quantizers = ccfg.num_quantizers + codec = AudioCodecDecoderSubmodule(codec=real, config=cfg) + label = AudioCodecDecoderSubmodule.DECODE_REGION + region = _captured_on_cpu(codec.get_piecewise_accelerator_graph_configs(torch.device("cpu"), torch.float32)[label], + label) + rows, run = [], region.run + region.run = lambda *a, real_bs=None, **k: rows.append(real_bs) or run(*a, real_bs=real_bs, **k) + calls = {rid: torch.randint(0, ccfg.codebook_size, (23, ccfg.num_quantizers)) for rid in ("a", "b", "c")} + pieces = {rid: [] for rid in calls} + with torch.no_grad(): + starts = list(range(0, 23, 5)) + for start in starts: + inputs = [codec.prepare_inputs("codec_chunk", None, {"codec_tokens": [codes[start:start + 5]]}, + is_final_stream_chunk=start == starts[-1]) for codes in calls.values()] + eng = SimpleNamespace(request_ids=list(calls), resources={}, per_request_states=None, + piecewise_runners={label: region}) + out = codec.forward_batched("codec_chunk", eng, **codec.preprocess("codec_chunk", eng, inputs)) + for rid in calls: + pieces[rid].append(out[rid]["audio_chunk"][0]) + assert 3 in rows + for rid, codes in calls.items(): + whole, _ = real.decode(codes.unsqueeze(0)) + whole = (whole[0, 0].clamp(-1, 1) * 32767).to(torch.int16) + torch.testing.assert_close(torch.cat(pieces[rid]), whole, atol=1, rtol=0) + + +def _live_tensors() -> int: + import gc + + gc.collect() + return sum(type(o) is torch.Tensor for o in gc.get_objects()) + + +def test_codec_inverse_stft_keeps_nothing_per_length(): + """The standalone stream re-decodes the whole code history every frame, so + each frame is a new length: the inverse STFT must hold nothing per length.""" + from mstar.model.nemotron_duplex.components.audio_codec import _spec_to_wav + + n_fft, hop = 16, 4 + _spec_to_wav(torch.randn(1, n_fft // 2 + 1, 3, dtype=torch.complex64), n_fft, hop, n_fft) + before = _live_tensors() + for frames in range(4, 40): + _spec_to_wav(torch.randn(1, n_fft // 2 + 1, frames, dtype=torch.complex64), n_fft, hop, n_fft) + assert _live_tensors() == before + + +def test_codec_refuses_a_window_whose_overlap_add_has_a_gap(): + """The inverse STFT divides by the overlap-added squared window and zeroes + it where that is ~0; a window and hop that leave a gap are refused when the + codec is built, not met as silence mid-call.""" + from mstar.model.nemotron_duplex.components.audio_codec import check_window_envelope + + check_window_envelope(torch.hann_window(16), 4) + with pytest.raises(ValueError, match="overlap-add"): + check_window_envelope(torch.tensor([0.0, 0, 0, 1, 1, 0, 0, 0]), 4) + + def test_codec_requests_are_isolated(): codec = _make_codec() _run_codec(codec, "a", 5) @@ -453,14 +593,14 @@ def test_nano_frame_batch_fuses_rows_exactly_like_the_per_request_path(): torch.testing.assert_close(pre["input_embeds"], per_request) -def _small_perception(): +def _small_perception(**featurizer): """A real Perception, shrunk, with random weights: the streamer runs its modules.""" from mstar.model.nemotron_duplex.components.conformer import Perception from mstar.model.nemotron_duplex.config import ConformerSTTConfig torch.manual_seed(0) stt = ConformerSTTConfig(d_model=64, num_heads=4, num_layers=3, n_mels=32, feat_in=32, - subsampling_conv_channels=8, output_dim=48, att_context_size=(6, 0)) + subsampling_conv_channels=8, output_dim=48, att_context_size=(6, 0), **featurizer) p = Perception(stt) with torch.no_grad(): for t in p.parameters(): @@ -470,40 +610,115 @@ def _small_perception(): return p.eval() -def _make_encoder(max_session_samples=16000 * 10) -> ConformerEncoderSubmodule: - return ConformerEncoderSubmodule(_small_perception(), nn.Identity(), nn.Identity(), - config=NemotronDuplexConfig(), max_session_samples=max_session_samples) +class _Encoder: + """The encoder node with its resources, driven as the engine drives an eager + step: declare, admit and plan them, preprocess, forward, commit. The bounded + KV runs its torch path on CPU, in fp32 (the engine's KV dtype here).""" + + def __init__(self, **featurizer): + p = _small_perception(**featurizer) + self.node = ConformerEncoderSubmodule(p, nn.Identity(), nn.Identity(), config=NemotronDuplexConfig()) + def pool(key, num_layers, blocks): + spec = RecurrentStateSpec(key, {"conformer_encoder"}, + RecurrentStateConfig(num_layers=num_layers, blocks=blocks, max_slots=8)) + return RecurrentStatePool.build(spec, EngineResourceInfo(device=torch.device("cpu"))) + + self.pool = pool(ENCODER_STATE, p.config.num_layers, encoder_state_blocks(p.config, torch.float32)) + self.front = pool(ENCODER_FRONT, 1, front_state_blocks(p.config)) + self.kv = BoundedKVManager.build( + BoundedKVSpec(ENCODER_KV, {"conformer_encoder"}, encoder_kv_config(p.config, max_slots=8)), + EngineResourceInfo(device=torch.device("cpu"), kv_dtype=torch.float32), + ) + self.resources = {ENCODER_STATE: self.pool, ENCODER_FRONT: self.front, ENCODER_KV: self.kv} + self.runner = StepRunner(self.resources) + self.regions = {} # the piecewise runners, captured on CPU, by label, when a test sets them + + def prepare(self, handle, chunk, final=False): + inputs = {"audio_in": [chunk]} if chunk is not None else {"audio_in": []} + return self.node.prepare_inputs("encode", SimpleNamespace(rid_handle=handle), inputs, + is_final_stream_chunk=final) + + def declare(self, rids, inputs): + """The node's declaration, with the leases the engine takes for it first.""" + from mstar.engine.engine import Engine + + engine = Engine.__new__(Engine) + engine._submodules = {"conformer_encoder": SimpleNamespace(piecewise_runners=self.regions)} + ctx = StepContext(request_ids=tuple(rids), graph_walk="encode", slot=0, capture=False) + engine._maybe_lease_piecewise_regions("conformer_encoder", ctx, inputs) + return self.node.declare_step("encode", rids, inputs, piecewise_leases=ctx.piecewise_leases) + + def step(self, rids, inputs) -> BatchedModelOutput: + """One engine step over ``rids``: batched when the node can batch them, + else each call's own step, as the engine's per-request path runs it.""" + if len(rids) > 1 and not self.node.can_batch(None, inputs): + out = BatchedModelOutput(per_rid_outputs={}) + for rid, inp in zip(rids, inputs, strict=True): + out.per_rid_outputs.update(self._step([rid], [inp]).per_rid_outputs) + return out + return self._step(rids, inputs) + + def _step(self, rids, inputs) -> BatchedModelOutput: + for rid in rids: + for resource in self.resources.values(): + resource.ingest_request(rid) + step = self.declare(rids, inputs) + if step is not None: + step.set_ctx(StepContext(request_ids=tuple(rids), graph_walk="encode", slot=0, capture=False)) + assert self.runner.admit(step).ok + self.runner.plan(step) + engine_inputs = ModelInputsFromEngine( + request_ids=list(rids), per_request_info={}, resources=dict(self.resources), + piecewise_runners=dict(self.regions), step=step, + ) + with torch.no_grad(): + out = self.node.forward_batched("encode", engine_inputs, + **self.node.preprocess("encode", engine_inputs, inputs)) + if step is not None: + self.runner.commit(step) + return out + def encode(self, handle, chunk, final=False): + """One step for one call: what the stream hands the node, then its frames.""" + return self.step([handle], [self.prepare(handle, chunk, final)]).per_rid_outputs[handle]["audio_frame"] -def _prepare(enc, handle, chunk, final=False): - inputs = {"audio_in": [chunk]} if chunk is not None else {"audio_in": []} - return enc.prepare_inputs("encode", SimpleNamespace(rid_handle=handle), inputs, is_final_stream_chunk=final) + def offline(self, wav): + with torch.no_grad(): + return self.node.perception(wav.unsqueeze(0), torch.tensor([wav.numel()]))[0][0] -def _encode(enc, handle, chunk, final=False): - """One encoder step for one call, as the engine runs one unbatched: - prepare_inputs, preprocess, forward.""" - ni = _prepare(enc, handle, chunk, final) - eng = SimpleNamespace(request_ids=[handle]) - with torch.no_grad(): - out = enc.forward("encode", eng, **enc.preprocess("encode", eng, [ni])) - return out["audio_frame"] +def _encoder_regions(enc: _Encoder, *labels) -> dict: + """The encoder's piecewise regions as the engine runs them, captured on CPU.""" + configs = enc.node.get_piecewise_accelerator_graph_configs(torch.device("cpu"), torch.float32) + return {label: _captured_on_cpu(configs[label], label, enc.resources, enc.runner) for label in labels} def _encode_batch(enc, steps): """One batched encoder step, as the engine runs it: each ``(handle, chunk, - final)`` row through prepare_inputs, then one preprocess and forward_batched - over all of them. Returns each row's frames by handle.""" - nis = [_prepare(enc, h, c, f) for h, c, f in steps] - eng = SimpleNamespace(request_ids=[h for h, _c, _f in steps]) - with torch.no_grad(): - out = enc.forward_batched("encode", eng, **enc.preprocess("encode", eng, nis)) - return {h: out.per_rid_outputs[h]["audio_frame"] for h, _c, _f in steps} - - -def _offline(enc, wav): - with torch.no_grad(): - return enc.perception(wav.unsqueeze(0), torch.tensor([wav.numel()]))[0][0] + final)`` row through prepare_inputs, then one step over all of them. + Returns each row's frames by handle.""" + out = enc.step([h for h, _c, _f in steps], [enc.prepare(h, c, f) for h, c, f in steps]).per_rid_outputs + return {h: out[h]["audio_frame"] for h, _c, _f in steps} + + +def _stream(enc, wavs, sizes): + """Stream each call in its own chunk size, every call with audio left in + each step, then end them all; the frames each call got.""" + got = {h: [] for h in wavs} + pos = dict.fromkeys(wavs, 0) + while any(pos[h] < wavs[h].numel() for h in wavs): + rids, inputs = [], [] + for h in wavs: + if pos[h] < wavs[h].numel(): + inputs.append(enc.prepare(h, wavs[h][pos[h]:pos[h] + sizes[h]])) + pos[h] += sizes[h] + rids.append(h) + out = enc.step(rids, inputs).per_rid_outputs + for h in rids: + got[h] += out[h]["audio_frame"] + for h in wavs: + got[h] += enc.encode(h, None, final=True) + return got @pytest.mark.parametrize("end_rides_last_chunk", [False, True]) @@ -513,52 +728,197 @@ def test_encoder_streams_the_offline_frames(n, end_rides_last_chunk): on an empty final chunk or on the last chunk of audio (as an upload's usually is): the frames emitted in order are the offline encoding of the call, and the mel frames the offline featurizer's, the reflected right edge included - (a length on the hop grid is the one whose last frame reads it).""" - enc = _make_encoder() + (a length on the hop grid is the one whose last frame reads it), the newest + of which the front end holds.""" + enc = _Encoder() wav = torch.randn(n) * 0.1 - want = _offline(enc, wav) + want = enc.offline(wav) with torch.no_grad(): - mel = enc.perception.preprocessor(wav.unsqueeze(0), torch.tensor([n]))[0][0].T + mel = enc.node.perception.preprocessor(wav.unsqueeze(0), torch.tensor([n]))[0][0].T for handle, size in enumerate((320, 1280, 2560, 9000, 64000)): got = [] for at in range(0, wav.numel(), size): last = at + size >= wav.numel() - got += _encode(enc, handle, wav[at:at + size], final=last and end_rides_last_chunk) + got += enc.encode(handle, wav[at:at + size], final=last and end_rides_last_chunk) if not end_rides_last_chunk: - got += _encode(enc, handle, None, final=True) + got += enc.encode(handle, None, final=True) assert len(got) == want.shape[0] assert torch.allclose(torch.stack(got), want, atol=1e-5), size - st = enc.request_state(handle).get(enc.STATE_KEY) + st = enc.node.request_state(handle).get(enc.node.STATE_KEY) assert st.mels == mel.shape[0] - assert torch.allclose(st.mel[:st.mels], mel, atol=1e-5), size + held = enc.front.block("mel", 0)[enc.front._slots[handle]["main"].index] # the newest, right-aligned + assert torch.allclose(held, mel[-held.shape[0]:], atol=1e-5), size def test_encoder_batches_calls_at_different_points(): - """Two calls in one step, each with its own chunking, each get their own - offline frames.""" - enc = _make_encoder() - a, b = torch.randn(16000 * 2) * 0.1, torch.randn(16000 * 3) * 0.1 + """Calls in one step, each with its own chunking (one frame a step, several, + or none), each get their own offline frames.""" + enc = _Encoder() + wavs = {1: torch.randn(16000 * 2) * 0.1, 2: torch.randn(16000 * 3) * 0.1, 3: torch.randn(9000) * 0.1} + got = _stream(enc, wavs, {1: 1280, 2: 4000, 3: 320}) + for h, wav in wavs.items(): + assert torch.allclose(torch.stack(got[h]), enc.offline(wav), atol=1e-5), h + + +def test_encoder_live_steps_replay_the_regions_with_padding_rows(): + """Calls one frame of audio further on take the whole-step region; the + others in the same step (a call's first frames, one under a frame a step) + the eager front end and the layer region, both in buckets wider than the + batch. Every call gets its offline frames: the padding rows, which address + the pools' sinks, and the rows another path served leave them alone.""" + enc = _Encoder() + enc.regions = _encoder_regions(enc, ConformerEncoderSubmodule.STEP_REGION, ConformerEncoderSubmodule.LAYERS_REGION) + step, layers = enc.regions.values() + wavs = {1: torch.randn(16000 * 2) * 0.1, 2: torch.randn(16000) * 0.1, 3: torch.randn(6000) * 0.1} + got = _stream(enc, wavs, {1: 1280, 2: 1280, 3: 640}) + assert step.calls >= 20 and layers.calls >= 5 + for h, wav in wavs.items(): + assert torch.allclose(torch.stack(got[h]), enc.offline(wav), atol=1e-5), h + + +def test_encoder_layer_region_runs_only_the_calls_with_a_frame(): + """Calls one frame on take the whole-step region; the layer region then runs + only the other calls that got a frame, not every call of the step.""" + enc = _Encoder() + enc.regions = _encoder_regions(enc, ConformerEncoderSubmodule.STEP_REGION, ConformerEncoderSubmodule.LAYERS_REGION) + layers = enc.regions[ConformerEncoderSubmodule.LAYERS_REGION] + rows, run, step = [], layers.run, enc.step + layers.run = lambda *a, request_ids, **k: rows.append(len(request_ids)) or run(*a, request_ids=request_ids, **k) + enc.step = lambda rids, inputs: rows.append(f"step of {len(rids)}") or step(rids, inputs) + wavs = {1: torch.randn(16000) * 0.1, 2: torch.randn(16000) * 0.1, 3: torch.randn(16000) * 0.1} + got = _stream(enc, wavs, {1: 1280, 2: 1280, 3: 640}) + # once calls 1 and 2 are one frame on, the layer region runs call 3 alone + assert rows[4:10] == ["step of 3", 1, "step of 3", "step of 3", 1, "step of 3"], rows + for h, wav in wavs.items(): + assert torch.allclose(torch.stack(got[h]), enc.offline(wav), atol=1e-5), h + + +def test_encoder_paths_read_the_plan_current_when_they_use_it(monkeypatch): + """The KV attends under a plan of exactly the rows it is given, padding + included, with the eager path on this step's plan before a region plans its + own (the pools' gather / scatter_ read the plan the pool holds when they + run). A call with a chunk over a second runs in a step of its own.""" + enc = _Encoder() + enc.regions = _encoder_regions(enc, ConformerEncoderSubmodule.STEP_REGION, ConformerEncoderSubmodule.LAYERS_REGION) + attend = enc.kv.attend + + def attend_checked(layer_idx, q, *args, **kwargs): + assert len(enc.kv.current.rows) == q.shape[0], "the KV attended under a plan for other rows" + return attend(layer_idx, q, *args, **kwargs) + + monkeypatch.setattr(enc.kv, "attend", attend_checked) + wavs = {1: torch.randn(16000 * 2) * 0.1, 2: torch.randn(16000) * 0.1, 3: torch.randn(6000) * 0.1, + 4: torch.randn(16000 * 3) * 0.1} + got = _stream(enc, wavs, {1: 1280, 2: 1280, 3: 640, 4: 20000}) + assert all(region.calls for region in enc.regions.values()) + for h, wav in wavs.items(): + assert torch.allclose(torch.stack(got[h]), enc.offline(wav), atol=1e-5), h + + +def _kv_spans(step) -> list[int]: + return [seg.span for seg in step.get(ENCODER_KV).segments] + + +def test_encoder_an_all_steady_step_leaves_its_resources_to_its_region(): + """When every call is one frame on and the whole-step region is captured + for the batch, the region's own step owns the resources: the node declares + nothing, so no second plan is built and thrown away. Otherwise it declares + every call in each, and in the KV its new frames at its position.""" + enc = _Encoder() + wavs = {1: torch.randn(1280 * 6) * 0.1, 2: torch.randn(1280 * 6) * 0.1} + for at in (0, 1280): + enc.step([1, 2], [enc.prepare(h, wavs[h][at:at + 1280]) for h in (1, 2)]) + nxt = [enc.prepare(h, wavs[h][2560:3840]) for h in (1, 2)] + step = enc.declare([1, 2], nxt) # no region captured + assert set(step.keys()) == {ENCODER_FRONT, ENCODER_STATE, ENCODER_KV} + assert _kv_spans(step) == [1, 1] + states = [enc.node.request_state(h).get(enc.node.STATE_KEY) for h in (1, 2)] + assert [step.get(ENCODER_KV).positions[h].written for h in (1, 2)] == [st.frames for st in states] + # the call's own position, built when its frames last moved, not one per declaration + assert all(step.get(ENCODER_KV).positions[h] is st.position for h, st in zip((1, 2), states, strict=True)) + enc.regions = _encoder_regions(enc, ConformerEncoderSubmodule.STEP_REGION) + assert enc.declare([1, 2], nxt) is None + for at in range(2560, 1280 * 6, 1280): + out = enc.step([1, 2], [enc.prepare(h, wavs[h][at:at + 1280]) for h in (1, 2)]).per_rid_outputs + assert all(len(out[h]["audio_frame"]) == 1 for h in (1, 2)) + assert enc.regions[ConformerEncoderSubmodule.STEP_REGION].calls == 4 + mixed = [enc.prepare(1, torch.randn(640) * 0.1), enc.prepare(2, torch.randn(1280) * 0.1)] + step = enc.declare([1, 2], mixed) + assert set(step.keys()) == {ENCODER_FRONT, ENCODER_STATE, ENCODER_KV} + assert _kv_spans(step) == [mixed[0].kwargs["n"], 1] + + +def test_encoder_a_step_mixes_a_steady_call_with_eager_chunks_past_the_window(): + """One step holds a call one frame on, a call whose chunk adds more frames + than the attention window (the eager layers, a step of T > 1 where each + frame sees only the window before it, which the call one frame on then + joins) and a call under a frame a step: each gets its offline frames.""" + enc = _Encoder() + enc.regions = _encoder_regions(enc, ConformerEncoderSubmodule.STEP_REGION, ConformerEncoderSubmodule.LAYERS_REGION) + window = enc.node.perception.config.att_context_size[0] + wavs = {1: torch.randn(16000 * 3) * 0.1, 2: torch.randn(16000 * 3) * 0.1, 3: torch.randn(16000) * 0.1} + sizes = {1: 1280, 2: 16000, 3: 640} + assert 16000 // 1280 > window + widths, front = [], enc.node.streamer.front + + def spy_front(*args): + enc_in = front(*args) + widths.append(0 if enc_in is None else enc_in.x.shape[1]) + return enc_in + + enc.node.streamer.front = spy_front + got = _stream(enc, wavs, sizes) + assert enc.regions[ConformerEncoderSubmodule.STEP_REGION].calls > 0 and max(widths) > window + for h, wav in wavs.items(): + assert torch.allclose(torch.stack(got[h]), enc.offline(wav), atol=1e-5), h + + +def test_encoder_a_steady_call_joins_the_eager_layers_of_its_step(): + """A call one frame on whose window has wrapped, in a step where another + call's chunk makes several frames: this step's KV plan, which the eager + layers run on, writes both calls' frames, so the steady call can't also + take the whole-step region (whose write of its frame would follow one from + padding into the ring entry the region reads).""" + enc = _Encoder() + sub = ConformerEncoderSubmodule + enc.regions = _encoder_regions(enc, sub.STEP_REGION, sub.LAYERS_REGION) + window = enc.node.perception.config.att_context_size[0] + at = 1280 * (window + 4) + wavs = {1: torch.randn(at + 1280 * 3) * 0.1, 2: torch.randn(9000) * 0.1} got = {1: [], 2: []} - pos = {1: 0, 2: 0} - size = {1: 1280, 2: 4000} - wav = {1: a, 2: b} - while any(pos[h] < wav[h].numel() for h in (1, 2)): - rids, chunks, finals = [], [], [] - for h in (1, 2): - if pos[h] < wav[h].numel(): - ni = _prepare(enc, h, wav[h][pos[h]:pos[h] + size[h]]) - pos[h] += size[h] - rids.append(h) - chunks.append(ni.tensor_inputs["chunk"]) - finals.append(False) - with torch.no_grad(): - out = enc.forward_batched("encode", SimpleNamespace(request_ids=rids), chunk=chunks, final=finals) - for h in rids: - got[h] += out.per_rid_outputs[h]["audio_frame"] + for pos in range(0, at, 1280): + got[1] += enc.encode(1, wavs[1][pos:pos + 1280]) + out = _encode_batch(enc, [(1, wavs[1][at:at + 1280], False), (2, wavs[2], False)]) + assert len(out[2]) > 1 + for h in (1, 2): + got[h] += out[h] + got[1] += enc.encode(1, wavs[1][at + 1280:], final=True) + got[2] += enc.encode(2, None, final=True) for h in (1, 2): - got[h] += _encode(enc, h, None, final=True) - want = _offline(enc, wav[h]) - assert torch.allclose(torch.stack(got[h]), want, atol=1e-5) + assert torch.allclose(torch.stack(got[h]), enc.offline(wavs[h]), atol=1e-5), h + + +def test_encoder_runs_a_steady_call_eagerly_when_its_region_has_no_graph(): + """The whole-step region leases the batch's bucket, but replays only the + steady calls: when the runner dropped the capture for that sub-batch (here + one call's), they take the eager path, which this step's plans cover.""" + enc = _Encoder() + sub = ConformerEncoderSubmodule + enc.regions = _encoder_regions(enc, sub.STEP_REGION, sub.LAYERS_REGION) + step_region = enc.regions[sub.STEP_REGION] + step_region._graphs = {k: v for k, v in step_region._graphs.items() if k.bs != 1} + wavs = {1: torch.randn(1280 * 6) * 0.1, 2: torch.randn(1280 * 6) * 0.1} + got = {1: [], 2: []} + for at in (0, 1280, 2560): + for h, frames in _encode_batch(enc, [(h, wavs[h][at:at + 1280], False) for h in (1, 2)]).items(): + got[h] += frames + # call 1 is one frame on (steady), call 2 appends half a frame + out = _encode_batch(enc, [(1, wavs[1][3840:5120], False), (2, wavs[2][3840:4480], False)]) + for h in (1, 2): + got[h] += out[h] + for h, end in ((1, 5120), (2, 4480)): + got[h] += enc.encode(h, wavs[h][end:], final=True) + assert torch.allclose(torch.stack(got[h]), enc.offline(wavs[h]), atol=1e-5), h @pytest.mark.parametrize("b_end_rides", [False, True]) @@ -567,11 +927,10 @@ def test_encoder_batches_calls_at_different_points(): def test_encoder_a_calls_last_frames_do_not_depend_on_its_batch(b_len, a_chunk, b_end_rides): """Call B ends in the same step as call A's longer chunk, A still mid-call: B's last frames, which read past its audio, are still its offline ones.""" - enc = _make_encoder() + enc = _Encoder() starts = list(range(0, b_len, 1280)) steps = len(starts) + (0 if b_end_rides else 1) a, b = torch.randn(a_chunk * steps) * 0.1, torch.randn(b_len) * 0.1 - enc.max_session_samples = a.numel() got_b = [] for k, pb in enumerate(starts): b_final = b_end_rides and k == len(starts) - 1 @@ -579,39 +938,153 @@ def test_encoder_a_calls_last_frames_do_not_depend_on_its_batch(b_len, a_chunk, got_b += _encode_batch(enc, rows)[2] if not b_end_rides: got_b += _encode_batch(enc, [(1, a[-a_chunk:], False), (2, None, True)])[2] - want = _offline(enc, b) + want = enc.offline(b) assert len(got_b) == want.shape[0] assert torch.allclose(torch.stack(got_b), want, atol=1e-5) -def test_encoder_runs_a_chunk_over_a_second_in_a_pass_of_its_own(): - """A streamer pass pads every call to its longest chunk: an upload or a large - append runs alone, and only the calls with a second or less share a pass.""" - enc = _make_encoder() - enc.forward_batched("encode", SimpleNamespace(request_ids=[]), chunk=[], final=[]) # builds the streamer - steps = [] - step = enc._streamer.step - enc._streamer.step = lambda sts, chunks, final: steps.append( - sorted(0 if c is None else c.numel() for c in chunks)) or step(sts, chunks, final) +def test_encoder_keeps_a_chunk_over_a_second_out_of_a_batch(): + """A pass pads every call to its longest chunk, so a step holding an upload + or a large append is not batched: can_batch refuses it, by the chunk length + prepare_inputs reports, and the engine runs each call's step on its own. + Calls with a second or less share one pass. The frames are the offline + ones either way.""" + enc = _Encoder() + streamer, fronts = enc.node.streamer, [] + front = streamer.front + + def spy_front(sts, chunks, final, *rest): + fronts.append(sorted(c.numel() for c in chunks if c is not None)) + return front(sts, chunks, final, *rest) + + streamer.front = spy_front wav = {1: torch.randn(16000 * 5) * 0.1, 2: torch.randn(1280) * 0.1, 3: torch.randn(16001) * 0.1, 4: torch.randn(16000) * 0.1} - ni = [_prepare(enc, h, w, final=True) for h, w in wav.items()] - eng = SimpleNamespace(request_ids=list(wav)) - with torch.no_grad(): - out = enc.forward_batched("encode", eng, **enc.preprocess("encode", eng, ni)) - assert sorted(steps) == [[1280, 16000], [16001], [16000 * 5]] + prepared = {h: enc.prepare(h, w, final=True) for h, w in wav.items()} + assert {h: inp.input_seq_len for h, inp in prepared.items()} == {h: w.numel() for h, w in wav.items()} + assert not enc.node.can_batch(None, list(prepared.values())) + assert not enc.node.can_batch(None, [prepared[2], prepared[3]]) + assert enc.node.can_batch(None, [prepared[2], prepared[4]]) + assert enc.node.can_batch(None, [prepared[1]]) # alone, it is its own step + + out = _encode_batch(enc, [(h, w, True) for h, w in wav.items()]) + assert sorted(fronts) == [[1280], [16000], [16001], [16000 * 5]] # every call on its own for h, w in wav.items(): - assert torch.allclose(torch.stack(out.per_rid_outputs[h]["audio_frame"]), _offline(enc, w), atol=1e-5) + assert torch.allclose(torch.stack(out[h]), enc.offline(w), atol=1e-5) + + fronts.clear() + short = {5: torch.randn(1280) * 0.1, 6: torch.randn(16000) * 0.1} + out = _encode_batch(enc, [(h, w, True) for h, w in short.items()]) + assert fronts == [[1280, 16000]] # one pass + for h, w in short.items(): + assert torch.allclose(torch.stack(out[h]), enc.offline(w), atol=1e-5) + + +def test_encoder_and_codec_take_no_batch_wider_than_their_graphs(): + """Their graphs are piecewise regions, which the engine's batch cap does not + see: a wider step would run every call eagerly, so each node caps its batch + at its widest capture and the scheduler splits the rest.""" + from mstar.engine.engine import Engine + + nodes = { + "conformer_encoder": ConformerEncoderSubmodule(nn.Identity(), nn.Identity(), nn.Identity(), + config=NemotronDuplexConfig()), + "audio_codec": AudioCodecDecoderSubmodule(codec=nn.Identity(), config=NemotronDuplexConfig()), + } + engine = Engine.__new__(Engine) + engine._submodules = {name: SimpleNamespace(submodule=node, accelerator_graph_runner=None, + capture_runner=None) + for name, node in nodes.items()} + for (name, node), walk in zip(nodes.items(), ("encode", "codec_chunk"), strict=True): + assert engine.get_max_batch_size(name, walk) == max(node.CAPTURE_BATCH_SIZES), name + + +@pytest.mark.parametrize("n_fft", [512, 1024]) +def test_encoder_steady_steps_match_offline_for_the_featurizer_they_get(n_fft): + """The whole-step region's slices assume the featurizer's half window spans + more than one hop and at most two (512 / 160); with any other the encoder + serves every step through the general path, which takes any geometry.""" + enc = _Encoder(n_fft=n_fft) + regions = _encoder_regions(enc, ConformerEncoderSubmodule.STEP_REGION) if n_fft == 512 else {} + assert (ConformerEncoderSubmodule.STEP_REGION in enc.node.get_piecewise_accelerator_graph_configs( + torch.device("cpu"), torch.float32)) == (n_fft == 512) + enc.regions = regions + wav = torch.randn(16000 * 2) * 0.1 + got = torch.stack(_stream(enc, {1: wav}, {1: 1280})[1]) + assert torch.allclose(got, enc.offline(wav), atol=1e-5) + if regions: + assert regions[ConformerEncoderSubmodule.STEP_REGION].calls > 0 + + +@pytest.mark.parametrize("n_fft", [512, 256]) +@pytest.mark.parametrize("n", [1580, 1728]) +def test_encoder_ends_a_call_with_its_last_frame_when_the_end_adds_no_mel_frame(n_fft, n): + """With a half window under a hop, the end of input can make a last encoder + frame final without adding a mel frame: the call still gets it.""" + enc = _Encoder(n_fft=n_fft, window_size=n_fft / 16000) + wav = torch.randn(n) * 0.1 + got = [] + for at in range(0, n, 320): + got += enc.encode(1, wav[at:at + 320]) + got += enc.encode(1, None, final=True) + want = enc.offline(wav) + assert len(got) == want.shape[0] + assert torch.allclose(torch.stack(got), want, atol=1e-5) + + +def test_encoder_decides_a_calls_steadiness_once_per_step(monkeypatch): + """Whether a call takes the whole-step region is decided once, where its + inputs are prepared; the declaration and the forward read that decision.""" + enc = _Encoder() + enc.regions = _encoder_regions(enc, ConformerEncoderSubmodule.STEP_REGION) + wavs = {h: torch.randn(1280 * 4) * 0.1 for h in (1, 2)} + + def decided_already(*args): + raise AssertionError("steadiness decided again after prepare_inputs") + + for at in range(0, 1280 * 4, 1280): + inputs = [enc.prepare(h, wavs[h][at:at + 1280]) for h in wavs] + with monkeypatch.context() as m: + m.setattr(enc.node.streamer, "is_steady", decided_already) + enc.step(list(wavs), inputs) + assert enc.regions[ConformerEncoderSubmodule.STEP_REGION].calls == 2 + + +def test_encoder_moves_a_steps_chunks_to_the_device_in_one_copy(): + """Each call's chunk arrives on the host; one copy per call would each be a + blocking transfer from pageable memory, so the step copies them as one.""" + from torch.profiler import ProfilerActivity, profile + + chunks = [torch.randn(1280), None, torch.randn(640), torch.randn(1280)] + with profile(activities=[ProfilerActivity.CPU]) as prof: + moved = ConformerEncoderSubmodule._to_device(chunks, torch.device("meta")) + copies = [e for e in prof.events() if e.name == "aten::_to_copy"] + assert len(copies) == 1 + assert [None if c is None else (c.device.type, c.shape) for c in moved] == [ + ("meta", (1280,)), None, ("meta", (640,)), ("meta", (1280,))] + + +def test_encoder_step_region_replays_three_calls_with_a_padding_row(): + """Three calls one frame on replay the whole-step region in a bucket of + four: the padding row addresses both pools' sinks and the calls get their + offline frames.""" + enc = _Encoder() + enc.regions = _encoder_regions(enc, ConformerEncoderSubmodule.STEP_REGION) + region = enc.regions[ConformerEncoderSubmodule.STEP_REGION] + rows, run = [], region.run + region.run = lambda *a, request_ids, **k: rows.append(len(request_ids)) or run(*a, request_ids=request_ids, **k) + wavs = {h: torch.randn(16000) * 0.1 for h in (1, 2, 3)} + got = _stream(enc, wavs, dict.fromkeys(wavs, 1280)) + assert rows and set(rows) == {3} + for h, wav in wavs.items(): + assert torch.allclose(torch.stack(got[h]), enc.offline(wav), atol=1e-5), h def test_encoder_offers_nothing_to_the_stop_check(): """The encoder has no check_stop and its stream ends with the client's: its frames are not copied to the host every step for one.""" - enc = _make_encoder() - ni = _prepare(enc, 1, torch.randn(4000) * 0.1) - eng = SimpleNamespace(request_ids=[1]) - with torch.no_grad(): - out = enc.forward_batched("encode", eng, **enc.preprocess("encode", eng, [ni])) + enc = _Encoder() + out = enc.step([1], [enc.prepare(1, torch.randn(4000) * 0.1)]) assert isinstance(out, BatchedModelOutput) and out.check_stop_buffers == {} assert len(out.per_rid_outputs[1]["audio_frame"]) > 0 @@ -619,32 +1092,32 @@ def test_encoder_offers_nothing_to_the_stop_check(): def test_encoder_projects_the_relative_positions_once_per_layer(): """The positions a query can see are fixed, so each layer's projection of them is computed once, not every step.""" - enc = _make_encoder() + enc = _Encoder() calls = [] - for layer in enc.perception.encoder.layers: + for layer in enc.node.perception.encoder.layers: layer.self_attn.linear_pos.register_forward_hook(lambda *a: calls.append(1)) wav = torch.randn(16000) * 0.1 for at in range(0, wav.numel(), 1280): - _encode(enc, 1, wav[at:at + 1280]) - assert len(calls) == len(enc.perception.encoder.layers) + enc.encode(1, wav[at:at + 1280]) + assert len(calls) == len(enc.node.perception.encoder.layers) def test_encoder_waits_for_enough_audio_to_featurize(): - enc = _make_encoder() - assert _encode(enc, 1, torch.zeros(10)) == [] + enc = _Encoder() + assert enc.encode(1, torch.zeros(10)) == [] @pytest.mark.parametrize("n", [257, 1000, 1279, 1281]) def test_encoder_a_short_call_gets_the_offline_frames(n): """Offline encodes any call longer than n_fft/2 samples, under one frame too.""" - enc = _make_encoder() + enc = _Encoder() wav = torch.randn(n) * 0.1 - want = _offline(enc, wav) + want = enc.offline(wav) for handle, size in enumerate((100, n)): got = [] for at in range(0, n, size): - got += _encode(enc, handle, wav[at:at + size]) - got += _encode(enc, handle, None, final=True) + got += enc.encode(handle, wav[at:at + size]) + got += enc.encode(handle, None, final=True) assert len(got) == want.shape[0] assert torch.allclose(torch.stack(got), want, atol=1e-5), size @@ -653,16 +1126,136 @@ def test_encoder_a_short_call_gets_the_offline_frames(n): def test_encoder_fails_a_call_that_ends_without_enough_audio(chunks): """No audio, or no more than n_fft/2 samples (which offline can't featurize): the call fails rather than reply to nothing.""" - enc = _make_encoder() + enc = _Encoder() for chunk in chunks: - _encode(enc, 1, chunk) + enc.encode(1, chunk) with pytest.raises(ValueError, match="ended with"): - _encode(enc, 1, None, final=True) + enc.encode(1, None, final=True) def test_encoder_state_is_per_call_and_cleaned_up(): - enc = _make_encoder() - _encode(enc, 1, torch.randn(4000) * 0.1) - _encode(enc, 2, torch.randn(4000) * 0.1) - enc.cleanup_request(1) - assert 1 not in enc.request_states and 2 in enc.request_states + enc = _Encoder() + free = enc.pool.num_free_slots + enc.encode(1, torch.randn(4000) * 0.1) + enc.encode(2, torch.randn(4000) * 0.1) + assert enc.pool.num_free_slots == enc.front.num_free_slots == free - 2 + assert set(enc.kv._requests) == {1, 2} + enc.node.cleanup_request(1) + for resource in enc.resources.values(): + resource.remove_request(1) + assert 1 not in enc.node.request_states and 2 in enc.node.request_states + assert enc.pool.num_free_slots == enc.front.num_free_slots == free - 1 + assert set(enc.kv._requests) == {2} + + +# -- no cuFFT in a captured region --------------------------------------------- + + +def _without_fft(config): + """``config`` with its region raising on any FFT op: torch.stft and + torch.fft reach the dispatcher as ``aten._fft_*``.""" + from torch.utils._python_dispatch import TorchDispatchMode + + class _NoFFT(TorchDispatchMode): + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + assert not func.__name__.startswith("_fft_"), f"{func} in a captured region" + return func(*args, **(kwargs or {})) + + fn = config.capture_fn + + def region(call): + with _NoFFT(): + return fn(call) + + config.capture_fn = region + return config + + +def test_encoder_regions_run_no_fft(): + """A captured cuFFT plan lives in torch's plan cache, whose LRU eviction + destroys it under the graph; so the regions take their STFT as a matmul. + The eager path, which the guard leaves alone, keeps torch.stft.""" + enc = _Encoder() + configs = enc.node.get_piecewise_accelerator_graph_configs(torch.device("cpu"), torch.float32) + labels = (ConformerEncoderSubmodule.STEP_REGION, ConformerEncoderSubmodule.LAYERS_REGION) + enc.regions = {label: _captured_on_cpu(_without_fft(configs[label]), label, enc.resources, enc.runner) + for label in labels} + wavs = {1: torch.randn(16000) * 0.1, 2: torch.randn(16000 + 640) * 0.1} + got = _stream(enc, wavs, {1: 1280, 2: 640}) + assert all(enc.regions[label].calls for label in labels) + for h, wav in wavs.items(): + assert torch.allclose(torch.stack(got[h]), enc.offline(wav), atol=1e-5), h + + +def test_codec_region_runs_no_fft(): + """As the encoder's: the decode region's inverse DFT is a matmul, and its + stream is still the full-history decode.""" + from mstar.model.nemotron_duplex.components.audio_codec import AudioCodec + from mstar.model.nemotron_duplex.config import CodecConfig + + torch.manual_seed(0) + ccfg = CodecConfig(base_hidden_size=8, latent_size=16, codebook_size=16, num_quantizers=4) + real = AudioCodec(ccfg).eval() + for prm in real.parameters(): + if prm.is_floating_point(): + prm.data.normal_(0, 0.05) + cfg = NemotronDuplexConfig() + cfg.eartts.num_quantizers = ccfg.num_quantizers + codec = AudioCodecDecoderSubmodule(codec=real, config=cfg) + label = AudioCodecDecoderSubmodule.DECODE_REGION + region = _captured_on_cpu( + _without_fft(codec.get_piecewise_accelerator_graph_configs(torch.device("cpu"), torch.float32)[label]), + label) + codes = torch.randint(0, ccfg.codebook_size, (23, ccfg.num_quantizers)) + pieces = [] + with torch.no_grad(): + starts = list(range(0, 23, 5)) + for start in starts: + inp = codec.prepare_inputs("codec_chunk", None, {"codec_tokens": [codes[start:start + 5]]}, + is_final_stream_chunk=start == starts[-1]) + eng = SimpleNamespace(request_ids=["r"], resources={}, per_request_states=None, + piecewise_runners={label: region}) + pieces.append(codec.forward_batched("codec_chunk", eng, + **codec.preprocess("codec_chunk", eng, [inp]))["r"]["audio_chunk"][0]) + whole, _ = real.decode(codes.unsqueeze(0)) + assert region.calls >= 3 + whole = (whole[0, 0].clamp(-1, 1) * 32767).to(torch.int16) + torch.testing.assert_close(torch.cat(pieces), whole, atol=1, rtol=0) + + +@pytest.mark.parametrize("autocast", [False, True]) +@pytest.mark.parametrize(("n_fft", "win", "samples"), [(512, 400, 1632), (512, 400, 1700), (512, 512, 2000), + (256, 101, 900)]) +def test_stft_as_a_matmul_is_torch_stft(n_fft, win, samples, autocast): + """The matmul STFT is torch.stft's to fp32 rounding, measured against + float64; bf16 autocast leaves it in full precision.""" + from mstar.model.nemotron_duplex.components.dft import stft_as_real, stft_basis + + torch.manual_seed(0) + x, window = torch.randn(3, samples), torch.hann_window(win, periodic=False) + + def stft(signal): + return torch.view_as_real(torch.stft(signal, n_fft, hop_length=160, win_length=win, center=False, + window=window.to(signal.dtype), return_complex=True)) + + exact = stft(x.double()) + with torch.autocast("cpu", dtype=torch.bfloat16, enabled=autocast): + got = stft_as_real(x, stft_basis(window, n_fft), n_fft, 160) + assert got.dtype == torch.float32 and got.shape == exact.shape + fp32_fft = (stft(x) - exact).abs().max() + assert (got - exact).abs().max() <= fp32_fft + 1e-7 * exact.abs().max() + + +@pytest.mark.parametrize("n", [16, 15]) +def test_irfft_as_a_matmul_is_torch_irfft(n): + """The matmul inverse DFT is torch.fft.irfft's to fp32 rounding, measured + against float64, imaginary DC and Nyquist parts ignored as irfft ignores them.""" + from mstar.model.nemotron_duplex.components.dft import irfft_dim1 + + torch.manual_seed(0) + spec = torch.randn(4, n // 2 + 1, 50, dtype=torch.complex64) + exact = torch.fft.irfft(spec.to(torch.complex128), n=n, dim=-2) + got = irfft_dim1(spec, n) + assert got.dtype == torch.float32 and got.shape == exact.shape + fp32_fft = (torch.fft.irfft(spec, n=n, dim=-2) - exact).abs().max() + assert (got - exact).abs().max() <= fp32_fft + 1e-7 * exact.abs().max() diff --git a/test/modular/test_recurrent_gather_scatter.py b/test/modular/test_recurrent_gather_scatter.py index 36e64768e4..932429fd60 100644 --- a/test/modular/test_recurrent_gather_scatter.py +++ b/test/modular/test_recurrent_gather_scatter.py @@ -244,6 +244,23 @@ def test_a_pool_without_a_sink_refuses_gather_and_scatter(): pool.scatter_("h", rows_of(2, 0.0)) +def test_a_block_in_the_engine_dtype_takes_each_engines_own(): + """A ``dtype=None`` block is resolved by the pool, not written back into + the model's config, which its other specs and engines share.""" + config = RecurrentStateConfig( + num_layers=1, blocks={"h": RecurrentBlockConfig(shape=SHAPE, dtype=None)}, max_slots=MAX_SLOTS, + ) + pools = [ + RecurrentStatePool.build( + RecurrentStateSpec("state", {"node"}, config), + EngineResourceInfo(device=torch.device("cpu"), kv_dtype=kv_dtype), + ) + for kv_dtype in (torch.bfloat16, torch.float16) + ] + assert [pool.block("h", 0).dtype for pool in pools] == [torch.bfloat16, torch.float16] + assert config.blocks["h"].dtype is None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="captures a CUDA graph") def test_a_captured_graph_addresses_the_rows_each_replay_plans(): """Capture once over a bucket's addressing, then replay for batches of