From e993809be132bdf43ded1185179533119ba6923b Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 06:48:49 -0700 Subject: [PATCH 01/69] nemotron_duplex: capture the streaming encoder's layer stack A live step at 32 calls spent ~19 ms (in isolation; ~50 ms beside the worker's other threads) launching the encoder's ~1000 eager kernels, which put the live path past the 80 ms tick from 24 calls up. One-shot encodes a call once, so it never paid this. Each call's layer state (attention keys/values, conv inputs) moves into a RecurrentStatePool (ENCODER_STATE), and the layer stack runs over device counts, so a step where every call has at most one new frame (every live step) replays it as a piecewise CUDA graph. Steps with more frames (a one-shot clip, a large append, the call's end) run the same function eagerly. RecurrentStatePool gains gather/scatter_ (as in #383), and its addressing gains narrow: a chunk of more than a second still runs apart from the batch (no other call is padded to it), now on its own row of the step's plan. Also turns off per-message deflate on /generate/ws: PCM barely compresses, and it cost the event loop ~8% at 32 live calls. Co-Authored-By: Claude Opus 5.5 --- docs/models.rst | 4 + mstar/api_server/entrypoint.py | 7 +- mstar/engine/resources/recurrent/pool.py | 30 ++ .../nemotron_duplex/components/conformer.py | 178 ++++++----- mstar/model/nemotron_duplex/config.py | 1 + .../nemotron_duplex/nemotron_duplex_model.py | 14 +- mstar/model/nemotron_duplex/submodules.py | 118 ++++++- test/modular/test_nemotron_duplex_model.py | 17 +- .../test_nemotron_duplex_submodules.py | 287 ++++++++++++------ test/modular/test_recurrent_padding_rows.py | 21 ++ 10 files changed, 482 insertions(+), 195 deletions(-) diff --git a/docs/models.rst b/docs/models.rst index 8bf8b6be41..962b4c003d 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -514,6 +514,10 @@ 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 layer state (its attention keys / values and + conv inputs) in a recurrent-state pool, a slot per session; a live step, one + frame per call, replays its layer stack as a CUDA graph, and a step with more + frames (a one-shot upload, a large append) runs it eagerly. - 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/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index ba24d61ba2..700d3a9f26 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -2479,7 +2479,12 @@ def main(argv: list[str] | None = None): # 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) + uvicorn.Config( + app, host=args.host, port=args.port, access_log=False, + # /generate/ws streams PCM, which deflate barely shrinks; + # compressing it cost the event loop ~8% at 32 live calls + ws_per_message_deflate=False, + ) ) def _stop_server(): diff --git a/mstar/engine/resources/recurrent/pool.py b/mstar/engine/resources/recurrent/pool.py index 7a1786fbfe..3f1e0a68ac 100644 --- a/mstar/engine/resources/recurrent/pool.py +++ b/mstar/engine/resources/recurrent/pool.py @@ -65,6 +65,20 @@ class RecurrentAddressing: host_slot_indices: Sequence[int] | None = None host_has_state: Sequence[bool] | None = None + def narrow(self, start: int, length: int = 1) -> RecurrentAddressing: + """Rows ``start .. start + length`` of the step, for a forward that runs + them apart from the rest (a long row that would pad the others). Views + of this plan: valid exactly as long as it is.""" + end = start + length + assert 0 <= start < end <= self.num_rows, f"rows {start}..{end} of a step of {self.num_rows}" + return RecurrentAddressing( + slot_indices=self.slot_indices[start:end], + has_state=self.has_state[start:end], + num_rows=length, + host_slot_indices=None if self.host_slot_indices is None else self.host_slot_indices[start:end], + host_has_state=None if self.host_has_state is None else self.host_has_state[start:end], + ) + class RecurrentStatePool(Resource): @classmethod @@ -138,6 +152,22 @@ def block(self, name: str, layer_idx: int) -> torch.Tensor: """One layer's contiguous slot-major view of a block: ``[max_slots, *shape]``.""" return self._blocks[name][layer_idx] + def gather(self, name: str, rows: RecurrentAddressing, layer_idx: int = 0) -> torch.Tensor: + """Block ``name`` of these rows' slots, ``[rows, *shape]`` (a copy). + + ``rows`` is the step's addressing (``addressing``), so a captured and an + eager forward address their slots the same way; padding rows read the + sink.""" + index = rows.slot_indices[: rows.num_rows].long() + return self._blocks[name][layer_idx].index_select(0, index) + + def scatter_( + self, name: str, rows: RecurrentAddressing, value: torch.Tensor, layer_idx: int = 0, + ) -> None: + """Write ``value [rows, *shape]`` into block ``name`` of these rows' slots.""" + index = rows.slot_indices[: rows.num_rows].long() + self._blocks[name][layer_idx].index_copy_(0, index, value) + @property def num_free_slots(self) -> int: return len(self._free) diff --git a/mstar/model/nemotron_duplex/components/conformer.py b/mstar/model/nemotron_duplex/components/conformer.py index 920aa35980..3cfdf4c63d 100644 --- a/mstar/model/nemotron_duplex/components/conformer.py +++ b/mstar/model/nemotron_duplex/components/conformer.py @@ -19,14 +19,19 @@ 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.recurrent.config import RecurrentBlockConfig from mstar.model.nemotron_duplex.components._util import param from mstar.model.nemotron_duplex.config import ConformerSTTConfig +if TYPE_CHECKING: + from mstar.engine.resources.recurrent.pool import RecurrentAddressing, 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 +332,35 @@ 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. + """One call's place in its audio stream. 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. + holds every mel frame so far. The layers' state (attention keys / values, + depthwise-conv inputs) lives in the encoder's recurrent-state pool. """ - def __init__(self, config: ConformerSTTConfig, max_samples: int, device, dtype): + def __init__(self, config: ConformerSTTConfig, max_samples: int, device): 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) 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) + self.frames = 0 # encoder frames emitted (= frames in the layers' state) + + +def encoder_state_blocks(config: ConformerSTTConfig, dtype: torch.dtype) -> dict[str, RecurrentBlockConfig]: + """A call's layer state, per layer: the last ``att_context`` keys and values + of its attention and the last ``kernel - 1`` inputs of its depthwise conv.""" + ctx = config.att_context_size[0] + return { + "k": RecurrentBlockConfig(shape=(config.num_heads, ctx, config.head_dim), dtype=dtype), + "v": RecurrentBlockConfig(shape=(config.num_heads, ctx, config.head_dim), dtype=dtype), + "conv": RecurrentBlockConfig(shape=(config.d_model, config.conv_kernel_size - 1), dtype=dtype), + } def final_mel_frames(config: ConformerSTTConfig, samples: int, final: bool) -> int: @@ -371,6 +382,14 @@ def final_encoder_frames(config: ConformerSTTConfig, mels: int, samples: int, fi return (mels - 1) // config.subsampling_factor + 1 +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 frames before the step and new ones.""" + x: torch.Tensor + frames: list[int] + 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 @@ -384,7 +403,9 @@ class ConformerStreamer: 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. + the depthwise conv its cached last 8 inputs, both from the pool. ``layers`` + is all device work over device counters, so a one-frame step can replay + as a CUDA graph. """ def __init__(self, perception: "Perception"): @@ -461,20 +482,59 @@ def _subsample(self, sts: list[StreamState], first: list[int], last: list[int], 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))] + def front(self, sts: list[StreamState], chunks: list[torch.Tensor | None], + final: list[bool]) -> EncoderInputs | None: + """Append each call's chunk; the encoder-input frames it made final, or + None when no call has one.""" + 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] + if not live: + return None + 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 + n = [0] * len(sts) + for i, st, f in zip(live, L, fin, strict=True): + n[i] = max(0, final_encoder_frames(c, st.mels, st.samples, f) - st.frames) + ready = [i for i in live if n[i]] + if not ready: + return None + R = [sts[i] for i in ready] + first = [st.frames for st in R] + last = [st.frames + n[i] - 1 for st, i in zip(R, ready, strict=True)] + new = self._subsample(R, first, last, [st.mels for st in R]) + x = new[0].new_zeros(len(sts), max(n), new[0].shape[-1]) + for i, e in zip(ready, new, strict=True): + x[i, : n[i]] = e + return EncoderInputs(x=x, frames=[st.frames for st in sts], n=n) + + @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): + st.frames += n + # -- layers -------------------------------------------------------------- - def _window(self, sts: list[StreamState], n: list[int], t: int, device) -> tuple[torch.Tensor, torch.Tensor]: + def _mask(self, frames: torch.Tensor, n: torch.Tensor, t: int) -> 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 + then the ``t`` new ones (the first ``n`` real); the same for every layer.""" + ctx, device = self._ctx, frames.device 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) + frames, n = frames.view(-1, 1, 1), n.view(-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) + mask = outside | (frames - ctx + kj.view(1, 1, -1) < 0) | (kj.view(1, 1, -1) - ctx >= n) return mask.unsqueeze(1), dist.clamp(0, ctx) def _attend(self, layer: "ConformerLayer", x: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, @@ -513,67 +573,33 @@ 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, pool: "RecurrentStatePool", rows: "RecurrentAddressing", + frames: torch.Tensor, n: torch.Tensor) -> torch.Tensor: + """The layer stack over ``x (B, t, D)`` and the projection, reading and + writing each call's state in ``pool`` at ``rows``. ``frames`` / ``n`` are + device ``(B,)`` counts: frames before the step, new frames (``n <= t``; a + call with none keeps its state as it was).""" + b, t, _ = x.shape + ctx, keep = self._ctx, self.c.conv_kernel_size - 1 + mask, dist = self._mask(frames, n, t) + # each call keeps the newest ctx keys/values and kernel-1 conv inputs: + # positions n .. n + keep of each concatenation + shift = n.view(b, 1, 1, 1) + kv_idx = (shift + torch.arange(ctx, device=x.device).view(1, 1, ctx, 1)).expand( + b, self.c.num_heads, ctx, self.c.head_dim) + conv_idx = (shift.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): + k_cache, v_cache = pool.gather("k", rows, li), pool.gather("v", rows, li) + conv_cache = pool.gather("conv", rows, 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) + a, k, v = self._attend(layer, layer.norm_self_att(r), k_cache, v_cache, mask, dist) r = r + a - cv, tails = self._conv(layer, layer.norm_conv(r), c_all[:, li]) + 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_("k", rows, k.gather(2, kv_idx).to(k_cache.dtype), li) + pool.scatter_("v", rows, v.gather(2, kv_idx).to(v_cache.dtype), li) + pool.scatter_("conv", rows, tails.gather(2, conv_idx).to(conv_cache.dtype), li) + return self.p.proj(x) diff --git a/mstar/model/nemotron_duplex/config.py b/mstar/model/nemotron_duplex/config.py index 3fed506bdd..ab01a355ea 100644 --- a/mstar/model/nemotron_duplex/config.py +++ b/mstar/model/nemotron_duplex/config.py @@ -30,6 +30,7 @@ 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 layer state 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 238d58522e..a6059be743 100644 --- a/mstar/model/nemotron_duplex/nemotron_duplex_model.py +++ b/mstar/model/nemotron_duplex/nemotron_duplex_model.py @@ -69,8 +69,10 @@ from mstar.graph.special_destinations import EMIT_TO_CLIENT from mstar.model.base import ClientChunk, ForwardPassArgs, Model from mstar.model.multimodal import PromptPart +from mstar.model.nemotron_duplex.components.conformer import encoder_state_blocks from mstar.model.nemotron_duplex.config import ( AUDIO_IN, + ENCODER_STATE, MAMBA, MAMBA_STATE, NANO_ATTN, @@ -403,7 +405,8 @@ 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 recurrent-state pool holding + # each call's layer state (one slot per session). The codec declares nothing. # ------------------------------------------------------------------- # Sessions resident at once (one recurrent slot each; ~137 MB per slot in @@ -470,6 +473,15 @@ def get_node_resources(self) -> list[NodeResourceSpec]: num_qo_heads=nano.num_attention_heads, ) return [ + RecurrentStateSpec( + resource_key=ENCODER_STATE, nodes={"conformer_encoder"}, + config=RecurrentStateConfig( + num_layers=self.config.stt.num_layers, + # bf16, the engine's autocast dtype, which the layers compute the state in + blocks=encoder_state_blocks(self.config.stt, torch.bfloat16), + max_slots=self.DEFAULT_MAMBA_SLOTS + 1, + ), + ), KVSpec(resource_key=NANO_KV, nodes={"nano_llm"}, config=kv_config), AttentionSpec( resource_key=NANO_ATTN, nodes={"nano_llm"}, diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index 212d09612f..9070da5135 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -28,6 +28,10 @@ from mstar.engine.accelerator_graph_config import ( AcceleratorGraphConfig, BatchedAcceleratorGraphConfig, + PiecewiseAcceleratorGraphConfig, + PiecewiseBatchedConfig, + PiecewiseCallInputs, + PiecewiseCaptureShape, ) from mstar.engine.resources import ( AttentionStep, @@ -43,6 +47,7 @@ from mstar.model.nemotron_duplex.components.conformer import ConformerStreamer, StreamState from mstar.model.nemotron_duplex.config import ( AUDIO_IN, + ENCODER_STATE, MAMBA, MAMBA_STATE, NANO_ATTN, @@ -429,15 +434,22 @@ class ConformerEncoderSubmodule(NodeSubmodule): 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. + A call's layer state (attention keys / values, conv inputs) lives in the + ``ENCODER_STATE`` recurrent pool; its mel front end in its per-request + state. A live step gives every call at most one frame, and its layer stack + replays a CUDA graph (``LAYERS_REGION``); a step with more (a one-shot + request's whole clip, a large append, the call's end) runs it eagerly. + 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. + # Variable-length front end, eager; the layer stack is the captured region. disable_torch_compile = True + LAYERS_REGION = "conformer_layers" + CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32] + 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. @@ -488,38 +500,114 @@ def preprocess(self, graph_walk, engine_inputs, inputs): def can_batch(self, batch, model_inputs) -> bool: return True + @staticmethod + def _declare_layers(request_ids) -> SubmoduleStep: + return SubmoduleStep( + segments=[Segment(request_id=rid, label="main", span=1) for rid in request_ids], + steps={ENCODER_STATE: RecurrentStep()}, + ) + + def declare_step(self, graph_walk, request_ids, inputs, slot_lease=None, piecewise_leases=None, + **kwargs) -> SubmoduleStep: + # Always the eager path's addressing. When the step replays the region + # instead, the region admits, plans and commits the same slots under + # its own step; for this pool that only re-reads the slot table and + # marks the slots written twice, so the two need no coordination. + del graph_walk, inputs, slot_lease, piecewise_leases, kwargs + return self._declare_layers(request_ids) + + @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]: + """The layer stack for one frame per call, which is every live step.""" + del tp_world_size + d_model = self.perception.config.d_model + + def make_static_inputs(shape: PiecewiseCaptureShape) -> dict[str, torch.Tensor]: + return { + "x": torch.zeros(shape.bs, 1, d_model, dtype=autocast_dtype, device=device), + "frames": torch.zeros(shape.bs, dtype=torch.long, device=device), + "n": torch.zeros(shape.bs, dtype=torch.long, device=device), + } + + return { + self.LAYERS_REGION: PiecewiseBatchedConfig( + capture_fn=self._layers_region, + make_static_inputs=make_static_inputs, + declare_step=lambda request_ids, seq_lens: self._declare_layers(request_ids), + seq_len=1, + capture_batch_sizes=self.CAPTURE_BATCH_SIZES, + compile=False, + ), + } + + 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"], pool, pool.addressing("main"), s["frames"], s["n"])} + def _stream_state(self, rid, device) -> 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(self.perception.config, self.max_session_samples, device) 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) 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] + final = list(final) 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] 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 + pool = engine_inputs.resources[ENCODER_STATE] + # A step pads every call to its longest chunk, so a chunk of more than a + # second (an upload, a large append) runs alone, on its own row of the + # plan; the batched pass then leaves that call as it is. + for i, c in enumerate(chunks): + if c is not None and c.numel() > longest: + enc = self.streamer.front([states[i]], [c], [final[i]]) + if enc is not None: + frames[i] = self._eager_layers(enc, pool, pool.addressing("main").narrow(i), device)[0] + self.streamer.advance([states[i]], enc) + chunks[i], final[i] = None, False + enc = self.streamer.front(states, chunks, final) + if enc is not None: + runner = (engine_inputs.piecewise_runners or {}).get(self.LAYERS_REGION) + if runner is not None and enc.x.shape[1] == 1 and runner.can_run(len(rids)): + counts = torch.tensor([enc.frames, enc.n], dtype=torch.long) + y = runner.run(static_inputs={"x": enc.x, "frames": counts[0], "n": counts[1]}, + request_ids=rids)["y"] + y = [row[:n] for row, n in zip(y, enc.n, strict=True)] + else: + y = self._eager_layers(enc, pool, pool.addressing("main"), device) + self.streamer.advance(states, enc) + for i, n in enumerate(enc.n): + if n: + frames[i] = y[i] # 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 _eager_layers(self, enc, pool, rows, device) -> list[torch.Tensor]: + """The layer stack eagerly over ``enc`` at ``rows``: each call's new frames.""" + counts = torch.tensor([enc.frames, enc.n], dtype=torch.long).to(device, non_blocking=True) + y = self.streamer.layers(enc.x, pool, rows, counts[0], counts[1]) + return [row[:n] for row, n in zip(y, enc.n, strict=True)] + def forward(self, graph_walk, engine_inputs, chunk=None, final=None, **kwargs) -> NameToTensorList: # one call's step: preprocess's lists of one row rid = engine_inputs.request_ids[0] diff --git a/test/modular/test_nemotron_duplex_model.py b/test/modular/test_nemotron_duplex_model.py index 2b66995c6a..bbca124dc8 100644 --- a/test/modular/test_nemotron_duplex_model.py +++ b/test/modular/test_nemotron_duplex_model.py @@ -27,6 +27,7 @@ from mstar.engine.resources.recurrent import Mamba2Geometry, RecurrentStateSpec from mstar.model.base import ClientChunk from mstar.model.nemotron_duplex.config import ( + ENCODER_STATE, MAMBA, MAMBA_STATE, NANO_ATTN, @@ -71,13 +72,23 @@ 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 recurrent pool holding each call's layer + state, a slot per session. 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_STATE, 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 = {ENCODER_STATE: {"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 enc.config.blocks["k"].shape == (stt.num_heads, stt.att_context_size[0], stt.head_dim) + assert enc.config.blocks["conv"].shape == (stt.d_model, stt.conv_kernel_size - 1) + assert enc.config.usable_slots == model.DEFAULT_MAMBA_SLOTS eartts = model.config.eartts tkv = by_key[TALKER_KV] diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index a1165b574f..82de2c2739 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -11,10 +11,17 @@ import torch from torch import nn +from mstar.engine.accelerator_graph_config import PiecewiseCallInputs, PiecewiseCaptureShape from mstar.engine.resources import AttentionStep, KVStep, SamplerStep +from mstar.engine.resources.base import EngineResourceInfo 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 BucketKey, SlotLease, StepContext +from mstar.model.nemotron_duplex.components.conformer import encoder_state_blocks from mstar.model.nemotron_duplex.config import ( + ENCODER_STATE, MAMBA, MAMBA_STATE, NANO_ATTN, @@ -28,7 +35,7 @@ ConformerEncoderSubmodule, NemotronHLLMSubmodule, ) -from mstar.model.submodule_base import BatchedModelOutput +from mstar.model.submodule_base import BatchedModelOutput, ModelInputsFromEngine H = 8 @@ -470,40 +477,117 @@ 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 state pool, driven as the engine drives an + eager step: declare, admit and plan the pool, preprocess, forward, commit.""" + + def __init__(self, max_session_samples=16000 * 10): + p = _small_perception() + self.node = ConformerEncoderSubmodule(p, nn.Identity(), nn.Identity(), config=NemotronDuplexConfig(), + max_session_samples=max_session_samples) + spec = RecurrentStateSpec(ENCODER_STATE, {"conformer_encoder"}, RecurrentStateConfig( + num_layers=p.config.num_layers, blocks=encoder_state_blocks(p.config, torch.float32), max_slots=8)) + self.pool = RecurrentStatePool.build(spec, EngineResourceInfo(device=torch.device("cpu"))) + self.runner = StepRunner({ENCODER_STATE: self.pool}) + self.region = None # a stand-in for the piecewise runner, when a test sets one + + 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 step(self, rids, inputs) -> BatchedModelOutput: + for rid in rids: + self.pool.ingest_request(rid) + step = self.node.declare_step("encode", rids, inputs) + 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={ENCODER_STATE: self.pool}, + piecewise_runners={} if self.region is None else {ConformerEncoderSubmodule.LAYERS_REGION: self.region}, + ) + with torch.no_grad(): + out = self.node.forward_batched("encode", engine_inputs, + **self.node.preprocess("encode", engine_inputs, inputs)) + 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"] +class _EagerRegion: + """The engine's piecewise runner on CPU: the same bucket padding, region + declaration and admit / plan / commit around the region, with the captured + callable run eagerly.""" + def __init__(self, enc: _Encoder, bs: int): + self.enc, self.bs, self.calls = enc, bs, 0 + self.config = enc.node.get_piecewise_accelerator_graph_configs( + torch.device("cpu"), torch.float32)[ConformerEncoderSubmodule.LAYERS_REGION] + self.dummies = [f"__pad_{i}__" for i in range(bs)] + for rid in self.dummies: + enc.pool.ingest_request(rid) -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 can_run(self, batch_size): + return batch_size <= self.bs + + def run(self, static_inputs, request_ids): + self.calls += 1 + n = len(request_ids) + static = self.config.make_static_inputs(PiecewiseCaptureShape(bs=self.bs, seq_lens=[1] * self.bs, + total_tokens=self.bs)) + for name, value in static_inputs.items(): + static[name][:n].copy_(value) + ids = [*request_ids, *self.dummies[n:]] + step = self.config.declare_step(ids, [1] * self.bs) + ctx = StepContext(request_ids=tuple(request_ids), graph_walk="pw", slot=0, capture=False, + slot_lease=SlotLease(slot=0, bucket=BucketKey("pw", self.bs, self.bs))) + ctx.set_padded_rids(tuple(ids)) + step.set_ctx(ctx) + assert self.enc.runner.admit(step).ok + self.enc.runner.plan(step) + out = self.config.capture_fn(PiecewiseCallInputs( + static_inputs=static, + engine_inputs=ModelInputsFromEngine(request_ids=ids, per_request_info={}, + resources={ENCODER_STATE: self.enc.pool}), + )) + self.enc.runner.commit(step) + return {name: value[:n].clone() for name, value in out.items()} -def _offline(enc, wav): - with torch.no_grad(): - return enc.perception(wav.unsqueeze(0), torch.tensor([wav.numel()]))[0][0] +def _encode_batch(enc, steps): + """One batched encoder step, as the engine runs it: each ``(handle, chunk, + 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]) @@ -514,51 +598,47 @@ def test_encoder_streams_the_offline_frames(n, end_rides_last_chunk): 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() + 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 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 - 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 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) + """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_region_with_padding_rows(): + """One frame a call a step replays the layer region, in a bucket wider than + the batch: the calls get their offline frames and the padding rows, which + address the pool's sink, leave them alone. A call under a frame a step + rides the region with no new frame.""" + enc = _Encoder() + enc.region = _EagerRegion(enc, bs=4) + 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 enc.region.calls >= 25 + for h, wav in wavs.items(): + assert torch.allclose(torch.stack(got[h]), enc.offline(wav), atol=1e-5), h @pytest.mark.parametrize("b_end_rides", [False, True]) @@ -567,11 +647,11 @@ 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() + enc.node.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 +659,44 @@ 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) + """A pass pads every call to its longest chunk: an upload or a large append + runs alone, on its own row of the step's plan, and only the calls with a + second or less share a pass, whose layers then take at most a second of + frames per call.""" + enc = _Encoder() + streamer, fronts, layers = enc.node.streamer, [], [] + front, layer_stack = streamer.front, streamer.layers + + 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) + + def spy_layers(x, pool, rows, *rest): + layers.append((x.shape[0], x.shape[1], rows.num_rows)) + return layer_stack(x, pool, rows, *rest) + + streamer.front, streamer.layers = spy_front, spy_layers 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]] + out = _encode_batch(enc, [(h, w, True) for h, w in wav.items()]) + assert sorted(fronts) == [[1280, 16000], [16001], [16000 * 5]] + assert all(b == rows for b, _t, rows in layers) + assert max(t for b, t, _rows in layers if b > 1) == len(out[4]) # the widest call of a second or less 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) 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 +704,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 +738,20 @@ 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 == free - 2 + enc.node.cleanup_request(1) + enc.pool.remove_request(1) + assert 1 not in enc.node.request_states and 2 in enc.node.request_states + assert enc.pool.num_free_slots == free - 1 diff --git a/test/modular/test_recurrent_padding_rows.py b/test/modular/test_recurrent_padding_rows.py index 5d1636f511..ebbbc43a24 100644 --- a/test/modular/test_recurrent_padding_rows.py +++ b/test/modular/test_recurrent_padding_rows.py @@ -145,3 +145,24 @@ def test_a_real_request_padding_and_release_across_the_rotation(): assert pool.num_free_slots == 1 pool.remove_request("r") assert pool.num_free_slots == 2 + + +def test_narrow_addresses_some_rows_of_the_step(): + """``narrow`` is a run of the step's rows, as views of its plan: a forward + that runs one row apart reads and writes that row's slot and no other.""" + pool = build_pool(max_slots=6) + rids = ["a", "b", "c"] + for rid in rids: + pool.ingest_request(rid) + step = decode_step(rids) + ctx = StepContext(request_ids=rids, graph_walk="decode", slot=0, capture=False) + assert pool.admit(step, ctx).ok + rows = pool.plan(step, ctx)["main"] + one = rows.narrow(1) + assert one.num_rows == 1 and one.slot_indices.tolist() == rows.slot_indices[1:2].tolist() + assert list(one.host_slot_indices) == list(rows.host_slot_indices[1:2]) + before = pool.block("conv", 0).clone() + pool.scatter_("conv", one, torch.ones(1, 96, 3)) + changed = [i for i in range(6) if not torch.equal(pool.block("conv", 0)[i], before[i])] + assert changed == [pool._slots["b"]["main"].index] + torch.testing.assert_close(pool.gather("conv", rows)[1:2], pool.gather("conv", one)) From 7f919cb6fa51b5ccf96e8a84c87bfaa9a77f3eda Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Sat, 17 Oct 2026 21:23:53 -0700 Subject: [PATCH 02/69] nemotron_duplex: capture the codec's steady-state decode Live calls start at different times, so their 5-frame codec chunks come due on different frames: at 24 staggered calls the codec ran ~5x the steps of a lockstep load (168 vs 32), ~4 rows each, every one an eager decode. A steady-state window (3 context frames + one chunk) now replays a piecewise CUDA graph per batch bucket; other window lengths (a call's first chunk, the flush) stay eager. Co-Authored-By: Claude Opus 5.5 --- mstar/model/nemotron_duplex/submodules.py | 47 +++++++++++++++++-- .../test_nemotron_duplex_submodules.py | 40 ++++++++++++++-- 2 files changed, 78 insertions(+), 9 deletions(-) diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index 9070da5135..eeb200007f 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -861,13 +861,19 @@ 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; fp32 codec. disable_torch_compile = True disable_autocast = True + DECODE_REGION = "codec_decode" + CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32] + CONTEXT_KEY = "codec_ctx" TAIL_KEY = "codec_tail" # the last chunk's held-back samples, int16 @@ -908,6 +914,34 @@ def preprocess(self, graph_walk, engine_inputs, inputs): def can_batch(self, batch, model_inputs) -> bool: return True + @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]: + audio, _ = self.codec.decode(call.static_inputs["codes"]) + 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 @@ -935,11 +969,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 = (getattr(engine_inputs, "piecewise_runners", None) 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_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index 82de2c2739..110250baf5 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -283,7 +283,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) @@ -389,12 +389,37 @@ def test_codec_cleanup_clears_per_request_context(): assert "r" not in codec.request_states +class _EagerDecodeRegion: + """The engine's piecewise runner on CPU for the codec's decode region: the + batch padded to its bucket with zero codes, the captured callable run eagerly.""" + + def __init__(self, codec, bs: int = 4): + self.bs, self.calls = bs, 0 + self.config = codec.get_piecewise_accelerator_graph_configs(torch.device("cpu"), torch.float32)[ + AudioCodecDecoderSubmodule.DECODE_REGION] + + def can_run(self, batch_size): + return batch_size <= self.bs + + def run(self, static_inputs, real_bs): + self.calls += 1 + static = self.config.make_static_inputs(PiecewiseCaptureShape(bs=self.bs, seq_lens=[1] * self.bs, + total_tokens=self.bs)) + for name, value in static_inputs.items(): + static[name][:real_bs].copy_(value) + out = self.config.capture_fn(PiecewiseCallInputs( + static_inputs=static, engine_inputs=ModelInputsFromEngine(request_ids=[], per_request_info={}))) + return {name: value[:real_bs].clone() for name, value in out.items()} + + +@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 @@ -407,6 +432,10 @@ 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 + runners[AudioCodecDecoderSubmodule.DECODE_REGION] = _EagerDecodeRegion(codec) chunk = 5 codes = torch.randint(0, ccfg.codebook_size, (frames, ccfg.num_quantizers)) with torch.no_grad(): @@ -418,10 +447,13 @@ 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 From e3fba4362a6524c461ea86b321171479b667dd27 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Sat, 17 Oct 2026 21:45:44 -0700 Subject: [PATCH 03/69] nemotron_duplex codec: build the iSTFT window envelope once per shape _spec_to_wav rebuilt the overlap-added squared window and asserted it nonzero on every decode: a host sync per codec step, which stalls the worker's pipeline and fails CUDA-graph capture. It depends only on the frame count, so build and check it once per shape and reuse it. Co-Authored-By: Claude Opus 5.5 --- .../nemotron_duplex/components/audio_codec.py | 24 +++++++++++++++---- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/mstar/model/nemotron_duplex/components/audio_codec.py b/mstar/model/nemotron_duplex/components/audio_codec.py index 12f94e577c..939ba7fac8 100644 --- a/mstar/model/nemotron_duplex/components/audio_codec.py +++ b/mstar/model/nemotron_duplex/components/audio_codec.py @@ -179,14 +179,28 @@ def _spec_to_wav( ifft, output_size=(1, output_size), kernel_size=(1, win_length), stride=(1, hop_length), )[..., 0, 0, pad:-pad] + wav = wav / _window_envelope(window, T, hop_length, win_length, output_size, pad) + return wav + + +_ENVELOPES: dict[tuple, torch.Tensor] = {} + + +def _window_envelope(window: torch.Tensor, T: int, hop_length: int, win_length: int, output_size: int, + pad: int) -> torch.Tensor: + """The overlap-added squared window ``_spec_to_wav`` divides by. It depends only + on the frame count, so it is built and checked (a host sync) once per shape, + which also keeps the decode free of syncs for CUDA-graph capture.""" + key = (T, hop_length, win_length, window.device) + env = _ENVELOPES.get(key) + if env is None: window_sq = window.square().expand(T, -1).transpose(0, 1) - window_envelope = F.fold( + env = F.fold( window_sq, 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 + assert (env > 1e-11).all(), "Window envelope has zero values, cannot normalize." + _ENVELOPES[key] = env + return env class AudioCodec(nn.Module): From ad124b5ba43e3198a8f823691e2b31180ca0860a Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Sat, 17 Oct 2026 22:54:37 -0700 Subject: [PATCH 04/69] nemotron_duplex encoder: bounded front end in a pool, a live step as one graph The eager front end (append, mel, subsample) cost ~10 ms of host time per live step at 32 calls, per-call Python over whole-call buffers. A step only reads a bounded window of each call: its last frame of audio (every sample until its first mel frame) and its last 16 mel frames. Those now live in a second recurrent pool (ENCODER_FRONT), and a step where every call appends one frame of audio on the frame grid runs whole (front end and layers) as a piecewise CUDA graph (STEP_REGION). Other calls in the step take the same math eagerly, bitwise equal to the previous streamer in fp32. Host counts reach the device with non-blocking copies. Co-Authored-By: Claude Opus 5.5 --- .../nemotron_duplex/components/conformer.py | 247 +++++++++++------- mstar/model/nemotron_duplex/config.py | 1 + .../nemotron_duplex/nemotron_duplex_model.py | 15 +- mstar/model/nemotron_duplex/submodules.py | 114 +++++--- test/modular/test_nemotron_duplex_model.py | 11 +- .../test_nemotron_duplex_submodules.py | 65 +++-- 6 files changed, 288 insertions(+), 165 deletions(-) diff --git a/mstar/model/nemotron_duplex/components/conformer.py b/mstar/model/nemotron_duplex/components/conformer.py index 3cfdf4c63d..d37e690d3f 100644 --- a/mstar/model/nemotron_duplex/components/conformer.py +++ b/mstar/model/nemotron_duplex/components/conformer.py @@ -332,26 +332,31 @@ def load_weights(self, weights): # --------------------------------------------------------------------------- class StreamState: - """One call's place in its audio stream. Host counters drive every slice, - so a step needs no sync. + """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``).""" - ``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. The layers' state (attention keys / values, - depthwise-conv inputs) lives in the encoder's recurrent-state pool. - """ + __slots__ = ("samples", "mels", "frames") - def __init__(self, config: ConformerSTTConfig, max_samples: int, device): - 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) - 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 layers' state) +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) -> dict[str, RecurrentBlockConfig]: """A call's layer state, per layer: the last ``att_context`` keys and values of its attention and the last ``kernel - 1`` inputs of its depthwise conv.""" @@ -395,7 +400,7 @@ class ConformerStreamer: 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 @@ -403,9 +408,13 @@ class ConformerStreamer: 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, both from the pool. ``layers`` - is all device work over device counters, so a one-frame step can replay - as a CUDA graph. + the depthwise conv its cached last 8 inputs. + + A call's state is its host counters (``StreamState``) and two pool slots: + the front end's (``front_state_blocks``) and the layers' + (``encoder_state_blocks``). ``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"): @@ -413,6 +422,7 @@ 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 # 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) @@ -421,100 +431,137 @@ def __init__(self, perception: "Perception"): # -- 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, + def _mel(self, batch: torch.Tensor) -> torch.Tensor: + """``(B, samples)`` preemphasized windows -> ``(B, T, n_mels)`` log-mel frames.""" + c, feat = self.c, self.p.preprocessor.featurizer + spec = 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) spec = torch.view_as_real(spec) 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 = 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, 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))] - - def front(self, sts: list[StreamState], chunks: list[torch.Tensor | None], - final: list[bool]) -> EncoderInputs | None: + 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 (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, frames: torch.Tensor, front: "RecurrentStatePool", + front_rows: "RecurrentAddressing", pool: "RecurrentStatePool", + rows: "RecurrentAddressing") -> 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. + + 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_rows), front.gather("sig", front_rows) + mel_tail = front.gather("mel", front_rows) + 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]) + window = torch.cat([mel_tail[:, 1:], mel_new], dim=1) # (B, 23, n_mels) + new = self._subsample(window)[:, 2:3] + front.scatter_("sig", front_rows, y) + front.scatter_("last", front_rows, x[:, -1:]) + front.scatter_("mel", front_rows, torch.cat([mel_tail, mel_new], dim=1)[:, -2 * f:]) + return self.layers(new, pool, rows, frames, torch.ones_like(frames)) + + def front(self, sts: list[StreamState], chunks: list[torch.Tensor | None], final: list[bool], + front: "RecurrentStatePool", front_rows: "RecurrentAddressing") -> EncoderInputs | None: """Append each call's chunk; the encoder-input frames it made final, or - None when no call has one.""" - 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] - if not live: + 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 - 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 + last, sig = front.gather("last", front_rows), front.gather("sig", front_rows) + mel_tail = front.gather("mel", front_rows) + new_sig, new_last, new_mel = list(sig), list(last), list(mel_tail) n = [0] * len(sts) - for i, st, f in zip(live, L, fin, strict=True): - n[i] = max(0, final_encoder_frames(c, st.mels, st.samples, f) - st.frames) - ready = [i for i in live if n[i]] - if not ready: + frames_before = [st.frames for st in sts] + 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)) + 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])) + for row, (i, lo, hi) in enumerate(lo_hi): + st = sts[i] + held = min(st.mels, keep_mel) + mels = torch.cat([mel_tail[i, keep_mel - held:], mel[row, : hi - lo]]) # mels[0] is frame lo - held + st.mels = hi + new_mel[i] = F.pad(mels[-keep_mel:], (0, 0, max(0, keep_mel - mels.shape[0]), 0)) + n[i] = max(0, final_encoder_frames(c, st.mels, st.samples, final[i]) - st.frames) + if n[i]: + start = max(0, f * st.frames - 2 * f) # the subsampling window's first mel frame + mel_rows.append((i, mels[start - (lo - held):], (st.frames - start // f, n[i]))) + front.scatter_("sig", front_rows, torch.stack(new_sig)) + front.scatter_("last", front_rows, torch.stack(new_last)) + front.scatter_("mel", front_rows, torch.stack(new_mel)) + if not mel_rows: return None - R = [sts[i] for i in ready] - first = [st.frames for st in R] - last = [st.frames + n[i] - 1 for st, i in zip(R, ready, strict=True)] - new = self._subsample(R, first, last, [st.mels for st in R]) - x = new[0].new_zeros(len(sts), max(n), new[0].shape[-1]) - for i, e in zip(ready, new, strict=True): - x[i, : n[i]] = e - return EncoderInputs(x=x, frames=[st.frames for st in sts], n=n) + 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, frames=frames_before, n=n) @staticmethod def advance(sts: list[StreamState], enc: EncoderInputs) -> None: diff --git a/mstar/model/nemotron_duplex/config.py b/mstar/model/nemotron_duplex/config.py index ab01a355ea..8abaa61966 100644 --- a/mstar/model/nemotron_duplex/config.py +++ b/mstar/model/nemotron_duplex/config.py @@ -31,6 +31,7 @@ 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 layer state +ENCODER_FRONT = "encoder_front" # recurrent-state pool: each call's streaming-encoder audio and mel tails 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 a6059be743..4942bb6148 100644 --- a/mstar/model/nemotron_duplex/nemotron_duplex_model.py +++ b/mstar/model/nemotron_duplex/nemotron_duplex_model.py @@ -69,9 +69,10 @@ from mstar.graph.special_destinations import EMIT_TO_CLIENT from mstar.model.base import ClientChunk, ForwardPassArgs, Model from mstar.model.multimodal import PromptPart -from mstar.model.nemotron_duplex.components.conformer import encoder_state_blocks +from mstar.model.nemotron_duplex.components.conformer import encoder_state_blocks, front_state_blocks from mstar.model.nemotron_duplex.config import ( AUDIO_IN, + ENCODER_FRONT, ENCODER_STATE, MAMBA, MAMBA_STATE, @@ -405,8 +406,9 @@ 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. conformer_encoder: a recurrent-state pool holding - # each call's layer state (one slot per session). The codec declares nothing. + # passes at call time. conformer_encoder: two recurrent-state pools, each + # call's layer state 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 @@ -482,6 +484,13 @@ def get_node_resources(self) -> list[NodeResourceSpec]: max_slots=self.DEFAULT_MAMBA_SLOTS + 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_MAMBA_SLOTS + 1, + ), + ), KVSpec(resource_key=NANO_KV, nodes={"nano_llm"}, config=kv_config), AttentionSpec( resource_key=NANO_ATTN, nodes={"nano_llm"}, diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index eeb200007f..e60b456292 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -47,6 +47,7 @@ from mstar.model.nemotron_duplex.components.conformer import ConformerStreamer, StreamState from mstar.model.nemotron_duplex.config import ( AUDIO_IN, + ENCODER_FRONT, ENCODER_STATE, MAMBA, MAMBA_STATE, @@ -447,6 +448,7 @@ class ConformerEncoderSubmodule(NodeSubmodule): # Variable-length front end, eager; the layer stack is the captured region. disable_torch_compile = True + STEP_REGION = "conformer_step" LAYERS_REGION = "conformer_layers" CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32] @@ -501,20 +503,20 @@ def can_batch(self, batch, model_inputs) -> bool: return True @staticmethod - def _declare_layers(request_ids) -> SubmoduleStep: + def _declare(request_ids, keys) -> SubmoduleStep: return SubmoduleStep( segments=[Segment(request_id=rid, label="main", span=1) for rid in request_ids], - steps={ENCODER_STATE: RecurrentStep()}, + steps={key: RecurrentStep() for key in keys}, ) def declare_step(self, graph_walk, request_ids, inputs, slot_lease=None, piecewise_leases=None, **kwargs) -> SubmoduleStep: - # Always the eager path's addressing. When the step replays the region + # Always the eager path's addressing. When the step replays a region # instead, the region admits, plans and commits the same slots under - # its own step; for this pool that only re-reads the slot table and + # its own step; for these pools that only re-reads the slot table and # marks the slots written twice, so the two need no coordination. del graph_walk, inputs, slot_lease, piecewise_leases, kwargs - return self._declare_layers(request_ids) + return self._declare(request_ids, (ENCODER_FRONT, ENCODER_STATE)) @property def streamer(self) -> ConformerStreamer: @@ -525,71 +527,119 @@ def streamer(self) -> ConformerStreamer: def get_piecewise_accelerator_graph_configs( self, device: torch.device, autocast_dtype: torch.dtype, tp_world_size: int = 1, ) -> dict[str, PiecewiseAcceleratorGraphConfig]: - """The layer stack for one frame per call, which is every live step.""" + """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 - d_model = self.perception.config.d_model + stt = self.perception.config + frame_samples = stt.hop_length * stt.subsampling_factor - def make_static_inputs(shape: PiecewiseCaptureShape) -> dict[str, torch.Tensor]: + def step_inputs(shape: PiecewiseCaptureShape) -> dict[str, torch.Tensor]: + return { + "chunk": torch.zeros(shape.bs, frame_samples, dtype=torch.float32, device=device), + "frames": torch.zeros(shape.bs, dtype=torch.long, device=device), + } + + def layers_inputs(shape: PiecewiseCaptureShape) -> dict[str, torch.Tensor]: return { - "x": torch.zeros(shape.bs, 1, d_model, dtype=autocast_dtype, device=device), + "x": torch.zeros(shape.bs, 1, stt.d_model, dtype=autocast_dtype, device=device), "frames": torch.zeros(shape.bs, dtype=torch.long, device=device), "n": torch.zeros(shape.bs, dtype=torch.long, device=device), } + def region(fn, inputs, keys) -> PiecewiseBatchedConfig: + return PiecewiseBatchedConfig( + capture_fn=fn, make_static_inputs=inputs, + declare_step=lambda request_ids, seq_lens: self._declare(request_ids, keys), + seq_len=1, capture_batch_sizes=self.CAPTURE_BATCH_SIZES, compile=False, + ) + return { - self.LAYERS_REGION: PiecewiseBatchedConfig( - capture_fn=self._layers_region, - make_static_inputs=make_static_inputs, - declare_step=lambda request_ids, seq_lens: self._declare_layers(request_ids), - seq_len=1, - capture_batch_sizes=self.CAPTURE_BATCH_SIZES, - compile=False, - ), + self.STEP_REGION: region(self._step_region, step_inputs, (ENCODER_FRONT, ENCODER_STATE)), + self.LAYERS_REGION: region(self._layers_region, layers_inputs, (ENCODER_STATE,)), } + def _step_region(self, call: PiecewiseCallInputs) -> dict[str, torch.Tensor]: + front, pool = call.resources[ENCODER_FRONT], call.resources[ENCODER_STATE] + s = call.static_inputs + return {"y": self.streamer.steady_step(s["chunk"], s["frames"], front, front.addressing("main"), + pool, pool.addressing("main"))} + 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"], pool, pool.addressing("main"), s["frames"], s["n"])} - def _stream_state(self, rid, device) -> StreamState: + def _stream_state(self, rid) -> StreamState: state = self.request_state(rid) st = state.get(self.STATE_KEY) if st is None: - st = StreamState(self.perception.config, self.max_session_samples, device) + st = StreamState() state.add(self.STATE_KEY, st) return st + @staticmethod + def _counts(values, 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) + def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=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] + streamer = self.streamer + states = [self._stream_state(rid) for rid in rids] + chunks = [c.to(device, non_blocking=True) if c is not None else None for c in chunk] final = list(final) longest = int(self.BATCHED_CHUNK_S * self.perception.config.sample_rate) + front, pool = engine_inputs.resources[ENCODER_FRONT], engine_inputs.resources[ENCODER_STATE] + # this step's own addressing: a region's plan replaces the pools' current one + front_rows, pool_rows = front.addressing("main"), pool.addressing("main") + runners = engine_inputs.piecewise_runners or {} frames: list[torch.Tensor | None] = [None] * len(rids) - pool = engine_inputs.resources[ENCODER_STATE] - # A step pads every call to its longest chunk, so a chunk of more than a + + # calls one frame of audio further on the frame grid: the whole step as one graph + step_runner = runners.get(self.STEP_REGION) + steady = [i for i, (st, c, f) in enumerate(zip(states, chunks, final, strict=True)) + if streamer.is_steady(st, c, f)] if step_runner is not None else [] + if steady and step_runner.can_run(len(steady)): + y = step_runner.run( + static_inputs={"chunk": torch.stack([chunks[i] for i in steady]), + "frames": self._counts([states[i].frames for i in steady], device)}, + 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] + chunks[i] = None # consumed; the rest of the step leaves the call as it is + + # A pass pads every call to its longest chunk, so a chunk of more than a # second (an upload, a large append) runs alone, on its own row of the # plan; the batched pass then leaves that call as it is. for i, c in enumerate(chunks): if c is not None and c.numel() > longest: - enc = self.streamer.front([states[i]], [c], [final[i]]) + enc = streamer.front([states[i]], [c], [final[i]], front, front_rows.narrow(i)) if enc is not None: - frames[i] = self._eager_layers(enc, pool, pool.addressing("main").narrow(i), device)[0] - self.streamer.advance([states[i]], enc) + frames[i] = self._eager_layers(enc, pool, pool_rows.narrow(i), device)[0] + streamer.advance([states[i]], enc) chunks[i], final[i] = None, False - enc = self.streamer.front(states, chunks, final) + + # the rest: the same math eagerly, the layers as a graph when no call got more than a frame + enc = streamer.front(states, chunks, final, front, front_rows) if enc is not None: - runner = (engine_inputs.piecewise_runners or {}).get(self.LAYERS_REGION) + runner = runners.get(self.LAYERS_REGION) if runner is not None and enc.x.shape[1] == 1 and runner.can_run(len(rids)): - counts = torch.tensor([enc.frames, enc.n], dtype=torch.long) + counts = self._counts([enc.frames, enc.n], device) y = runner.run(static_inputs={"x": enc.x, "frames": counts[0], "n": counts[1]}, request_ids=rids)["y"] y = [row[:n] for row, n in zip(y, enc.n, strict=True)] else: - y = self._eager_layers(enc, pool, pool.addressing("main"), device) - self.streamer.advance(states, enc) + y = self._eager_layers(enc, pool, pool_rows, device) + streamer.advance(states, enc) for i, n in enumerate(enc.n): if n: frames[i] = y[i] @@ -604,7 +654,7 @@ def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, **k def _eager_layers(self, enc, pool, rows, device) -> list[torch.Tensor]: """The layer stack eagerly over ``enc`` at ``rows``: each call's new frames.""" - counts = torch.tensor([enc.frames, enc.n], dtype=torch.long).to(device, non_blocking=True) + counts = self._counts([enc.frames, enc.n], device) y = self.streamer.layers(enc.x, pool, rows, counts[0], counts[1]) return [row[:n] for row, n in zip(y, enc.n, strict=True)] diff --git a/test/modular/test_nemotron_duplex_model.py b/test/modular/test_nemotron_duplex_model.py index bbca124dc8..23b7cbd0fe 100644 --- a/test/modular/test_nemotron_duplex_model.py +++ b/test/modular/test_nemotron_duplex_model.py @@ -27,6 +27,7 @@ from mstar.engine.resources.recurrent import Mamba2Geometry, RecurrentStateSpec from mstar.model.base import ClientChunk from mstar.model.nemotron_duplex.config import ( + ENCODER_FRONT, ENCODER_STATE, MAMBA, MAMBA_STATE, @@ -72,15 +73,17 @@ 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 a recurrent pool holding each call's layer - state, a slot per session. The codec owns no engine resource.""" + and RoPE positions; the encoder two recurrent pools, each call's layer state + 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) == {ENCODER_STATE, NANO_KV, NANO_ATTN, MAMBA_STATE, MAMBA, NANO_SAMPLER, + assert set(by_key) == {ENCODER_FRONT, ENCODER_STATE, 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} - node = {ENCODER_STATE: {"conformer_encoder"}, **dict.fromkeys(nano_keys, {"nano_llm"})} + node = {ENCODER_STATE: {"conformer_encoder"}, ENCODER_FRONT: {"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 diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index 110250baf5..b43da76f8f 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -19,8 +19,9 @@ from mstar.engine.resources.recurrent.pool import RecurrentStatePool from mstar.engine.resources.runner import StepRunner from mstar.engine.resources.step import BucketKey, SlotLease, StepContext -from mstar.model.nemotron_duplex.components.conformer import encoder_state_blocks +from mstar.model.nemotron_duplex.components.conformer import encoder_state_blocks, front_state_blocks from mstar.model.nemotron_duplex.config import ( + ENCODER_FRONT, ENCODER_STATE, MAMBA, MAMBA_STATE, @@ -517,11 +518,16 @@ def __init__(self, max_session_samples=16000 * 10): p = _small_perception() self.node = ConformerEncoderSubmodule(p, nn.Identity(), nn.Identity(), config=NemotronDuplexConfig(), max_session_samples=max_session_samples) - spec = RecurrentStateSpec(ENCODER_STATE, {"conformer_encoder"}, RecurrentStateConfig( - num_layers=p.config.num_layers, blocks=encoder_state_blocks(p.config, torch.float32), max_slots=8)) - self.pool = RecurrentStatePool.build(spec, EngineResourceInfo(device=torch.device("cpu"))) - self.runner = StepRunner({ENCODER_STATE: self.pool}) - self.region = None # a stand-in for the piecewise runner, when a test sets one + 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.resources = {ENCODER_STATE: self.pool, ENCODER_FRONT: self.front} + self.runner = StepRunner(self.resources) + self.regions = {} # stand-ins for the piecewise runners, 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": []} @@ -531,13 +537,14 @@ def prepare(self, handle, chunk, final=False): def step(self, rids, inputs) -> BatchedModelOutput: for rid in rids: self.pool.ingest_request(rid) + self.front.ingest_request(rid) step = self.node.declare_step("encode", rids, inputs) 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={ENCODER_STATE: self.pool}, - piecewise_runners={} if self.region is None else {ConformerEncoderSubmodule.LAYERS_REGION: self.region}, + request_ids=list(rids), per_request_info={}, resources=dict(self.resources), + piecewise_runners=dict(self.regions), ) with torch.no_grad(): out = self.node.forward_batched("encode", engine_inputs, @@ -559,13 +566,13 @@ class _EagerRegion: declaration and admit / plan / commit around the region, with the captured callable run eagerly.""" - def __init__(self, enc: _Encoder, bs: int): + def __init__(self, enc: _Encoder, bs: int, label: str): self.enc, self.bs, self.calls = enc, bs, 0 - self.config = enc.node.get_piecewise_accelerator_graph_configs( - torch.device("cpu"), torch.float32)[ConformerEncoderSubmodule.LAYERS_REGION] - self.dummies = [f"__pad_{i}__" for i in range(bs)] + self.config = enc.node.get_piecewise_accelerator_graph_configs(torch.device("cpu"), torch.float32)[label] + self.dummies = [f"__pad_{label}_{i}__" for i in range(bs)] for rid in self.dummies: - enc.pool.ingest_request(rid) + for pool in enc.resources.values(): + pool.ingest_request(rid) def can_run(self, batch_size): return batch_size <= self.bs @@ -588,7 +595,7 @@ def run(self, static_inputs, request_ids): out = self.config.capture_fn(PiecewiseCallInputs( static_inputs=static, engine_inputs=ModelInputsFromEngine(request_ids=ids, per_request_info={}, - resources={ENCODER_STATE: self.enc.pool}), + resources=dict(self.enc.resources)), )) self.enc.runner.commit(step) return {name: value[:n].clone() for name, value in out.items()} @@ -629,7 +636,8 @@ 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).""" + (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 = enc.offline(wav) @@ -646,7 +654,8 @@ def test_encoder_streams_the_offline_frames(n, end_rides_last_chunk): assert torch.allclose(torch.stack(got), want, atol=1e-5), size 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(): @@ -659,16 +668,19 @@ def test_encoder_batches_calls_at_different_points(): assert torch.allclose(torch.stack(got[h]), enc.offline(wav), atol=1e-5), h -def test_encoder_live_steps_replay_the_region_with_padding_rows(): - """One frame a call a step replays the layer region, in a bucket wider than - the batch: the calls get their offline frames and the padding rows, which - address the pool's sink, leave them alone. A call under a frame a step - rides the region with no new frame.""" +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.region = _EagerRegion(enc, bs=4) + step = _EagerRegion(enc, bs=4, label=ConformerEncoderSubmodule.STEP_REGION) + layers = _EagerRegion(enc, bs=4, label=ConformerEncoderSubmodule.LAYERS_REGION) + enc.regions = {ConformerEncoderSubmodule.STEP_REGION: step, ConformerEncoderSubmodule.LAYERS_REGION: layers} 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 enc.region.calls >= 25 + 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 @@ -782,8 +794,9 @@ def test_encoder_state_is_per_call_and_cleaned_up(): 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 == free - 2 + assert enc.pool.num_free_slots == enc.front.num_free_slots == free - 2 enc.node.cleanup_request(1) - enc.pool.remove_request(1) + for pool in enc.resources.values(): + pool.remove_request(1) assert 1 not in enc.node.request_states and 2 in enc.node.request_states - assert enc.pool.num_free_slots == free - 1 + assert enc.pool.num_free_slots == enc.front.num_free_slots == free - 1 From a612053ba36566b870cccf8bb97a0a4db32c6834 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 03:52:00 -0700 Subject: [PATCH 05/69] R1.1: the codec's inverse STFT keeps no envelope per length _spec_to_wav cached the overlap-added window envelope per frame count in a module-global dict that was never evicted. The standalone stream (DuplexStream) re-decodes the whole code history every 80 ms frame, so every frame was a new length and a new cached tensor: ~2 GiB after a minute. Build the envelope every call and guard the division on the device, as Kokoro's inverse STFT does, so the decode stays free of host syncs and capturable. Whether the window and hop can overlap-add to ~0 at all is checked once, when the codec is built, for every frame count. Co-Authored-By: Claude Opus 5.5 --- .../nemotron_duplex/components/audio_codec.py | 47 ++++++++++++------- .../test_nemotron_duplex_submodules.py | 31 ++++++++++++ 2 files changed, 60 insertions(+), 18 deletions(-) diff --git a/mstar/model/nemotron_duplex/components/audio_codec.py b/mstar/model/nemotron_duplex/components/audio_codec.py index 939ba7fac8..b6f6848454 100644 --- a/mstar/model/nemotron_duplex/components/audio_codec.py +++ b/mstar/model/nemotron_duplex/components/audio_codec.py @@ -179,28 +179,38 @@ def _spec_to_wav( ifft, output_size=(1, output_size), kernel_size=(1, win_length), stride=(1, hop_length), )[..., 0, 0, pad:-pad] - wav = wav / _window_envelope(window, T, hop_length, win_length, output_size, pad) - return wav + # 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)) -_ENVELOPES: dict[tuple, torch.Tensor] = {} +_ENVELOPE_FLOOR = 1e-11 -def _window_envelope(window: torch.Tensor, T: int, hop_length: int, win_length: int, output_size: int, - pad: int) -> torch.Tensor: - """The overlap-added squared window ``_spec_to_wav`` divides by. It depends only - on the frame count, so it is built and checked (a host sync) once per shape, - which also keeps the decode free of syncs for CUDA-graph capture.""" - key = (T, hop_length, win_length, window.device) - env = _ENVELOPES.get(key) - if env is None: - window_sq = window.square().expand(T, -1).transpose(0, 1) - env = F.fold( - window_sq, output_size=(1, output_size), kernel_size=(1, win_length), stride=(1, hop_length), - ).squeeze()[pad:-pad] - assert (env > 1e-11).all(), "Window envelope has zero values, cannot normalize." - _ENVELOPES[key] = env - return env +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] + + +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): @@ -215,6 +225,7 @@ 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 + check_window_envelope(torch.hann_window(cfg.n_fft), cfg.hop_length) # ---- autoencoder halves ------------------------------------------------- def ae_encode(self, x: torch.Tensor) -> torch.Tensor: diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index b43da76f8f..ff954ea2a6 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -459,6 +459,37 @@ def test_codec_stream_is_the_full_history_decode_with_the_real_decoder(frames, r torch.testing.assert_close(stream, whole, atol=1, rtol=0) # int16: one quantization step +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) From 876bdcb3f36926a151cf723fda642bbac609afb9 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 03:54:24 -0700 Subject: [PATCH 06/69] R1.2: refuse a session cap a per-session resource cannot hold The encoder's two new pools were sized from DEFAULT_MAMBA_SLOTS, while the docs and configs told operators to raise mamba_state's max_slots with max_concurrent_requests. Doing so left the encoder at 32 slots: call 33's encoder admit was refused and retried, with no frames, until another call ended. The talker KV pool had the same gap. The model now checks, at startup (validate_config_yaml, as whisper does), that max_concurrent_requests sessions fit every per-session resource after the deployment's overrides: mamba_state, encoder_state, encoder_front and talker_kv, naming the one that does not. The docs and configs name all four. Co-Authored-By: Claude Opus 5.5 --- configs/nemotron_duplex.yaml | 7 ++- configs/nemotron_duplex_disagg.yaml | 7 ++- docs/models.rst | 10 ++++- .../nemotron_duplex/nemotron_duplex_model.py | 44 +++++++++++++++++-- test/modular/test_nemotron_duplex_model.py | 25 +++++++++++ 5 files changed, 84 insertions(+), 9 deletions(-) diff --git a/configs/nemotron_duplex.yaml b/configs/nemotron_duplex.yaml index 4a6c8097e5..297ad7b750 100644 --- a/configs/nemotron_duplex.yaml +++ b/configs/nemotron_duplex.yaml @@ -9,8 +9,11 @@ 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; the overflow queues. +# Sessions at once; the overflow queues. Each session holds a slot in +# resources mamba_state, encoder_state and encoder_front and its pages in +# talker_kv, all sized for 32 (DEFAULT_MAMBA_SLOTS); to raise the cap, raise +# their max_slots (cap + 1, for the sink) and talker_kv's max_num_pages with it. +# The server refuses to start with a cap one of them cannot hold. max_concurrent_requests: 32 # A live call's input files over its whole life; past it the call ends with a diff --git a/configs/nemotron_duplex_disagg.yaml b/configs/nemotron_duplex_disagg.yaml index aafb665d9e..99d8a58042 100644 --- a/configs/nemotron_duplex_disagg.yaml +++ b/configs/nemotron_duplex_disagg.yaml @@ -14,8 +14,11 @@ 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; the overflow queues. +# Sessions at once; the overflow queues. Each session holds a slot in +# resources mamba_state, encoder_state and encoder_front and its pages in +# talker_kv, all sized for 32 (DEFAULT_MAMBA_SLOTS); to raise the cap, raise +# their max_slots (cap + 1, for the sink) and talker_kv's max_num_pages with it. +# The server refuses to start with a cap one of them cannot hold. max_concurrent_requests: 32 # A live call's input files over its whole life; past it the call ends with a diff --git a/docs/models.rst b/docs/models.rst index 962b4c003d..ccab3c2208 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -437,10 +437,16 @@ 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`` / ``state_dtype``), 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 its pages in ``talker_kv`` + (``max_num_pages``). The shipped configs set ``max_concurrent_requests: 32`` + to match; to serve more, raise the cap and all four 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. - ``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 diff --git a/mstar/model/nemotron_duplex/nemotron_duplex_model.py b/mstar/model/nemotron_duplex/nemotron_duplex_model.py index 4942bb6148..cb2334f0c9 100644 --- a/mstar/model/nemotron_duplex/nemotron_duplex_model.py +++ b/mstar/model/nemotron_duplex/nemotron_duplex_model.py @@ -62,6 +62,7 @@ ResourceReqConfig, SamplerSpec, SamplingReqConfig, + apply_yaml_overrides, ) from mstar.engine.resources.linear_attn.config import LinearAttnConfig, LinearAttnSpec, LinearAttnVariant from mstar.engine.resources.recurrent import Mamba2Geometry, RecurrentStateConfig, RecurrentStateSpec @@ -448,11 +449,48 @@ 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.""" + def talker_kv_pages(self, page_size: int, sessions: int | None = None) -> int: + """Pages for ``sessions`` (DEFAULT_MAMBA_SLOTS) sessions of max_session_frames, both CFG streams.""" + sessions = self.DEFAULT_MAMBA_SLOTS if sessions is None else sessions positions = self.TALKER_WARMUP_POSITIONS + self.max_session_frames + 1 # + 1 for the pool's sink page - return self.DEFAULT_MAMBA_SLOTS * 2 * -(-positions // page_size) + 1 + return sessions * 2 * -(-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``) and its 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 spec.resource_key == TALKER_KV: + kv = spec.config + per_session = self.talker_kv_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 diff --git a/test/modular/test_nemotron_duplex_model.py b/test/modular/test_nemotron_duplex_model.py index 23b7cbd0fe..f45644c8b4 100644 --- a/test/modular/test_nemotron_duplex_model.py +++ b/test/modular/test_nemotron_duplex_model.py @@ -132,6 +132,31 @@ 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)) + + +@pytest.mark.parametrize("key", [MAMBA_STATE, ENCODER_STATE, ENCODER_FRONT, TALKER_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 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() + pages = 48 * (model.talker_kv_pages(128) - 1) // model.DEFAULT_MAMBA_SLOTS + 1 + resources = {MAMBA_STATE: {"max_slots": 49}, ENCODER_STATE: {"max_slots": 49}, + ENCODER_FRONT: {"max_slots": 49}, TALKER_KV: {"max_num_pages": pages}} + 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") + + def test_duplex_request_resource_configs(): model = _make_model() cfg = model.get_request_resource_configs({}, None) From 0732042580b20627aca78380cef6fc5a36bcfc7f Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 04:04:04 -0700 Subject: [PATCH 07/69] R1.16: drive the encoder's and codec's regions through the real runner The tests' hand-written stand-ins for the piecewise runner shared one bucket key across batch sizes and rebuilt their inputs every call, so they could not catch what the engine's runner does differently (per-bucket static buffers, padding rows, its own admit / plan / commit). The tests now capture the real PiecewiseCudaGraphRunner on CPU, with a graph whose replay re-runs the region. Co-Authored-By: Claude Opus 5.5 --- .../test_nemotron_duplex_submodules.py | 106 ++++++++---------- 1 file changed, 44 insertions(+), 62 deletions(-) diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index ff954ea2a6..cff1a8a57b 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -11,14 +11,13 @@ import torch from torch import nn -from mstar.engine.accelerator_graph_config import PiecewiseCallInputs, PiecewiseCaptureShape from mstar.engine.resources import AttentionStep, KVStep, SamplerStep from mstar.engine.resources.base import EngineResourceInfo from mstar.engine.resources.linear_attn.config import LinearAttnStep 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 BucketKey, SlotLease, StepContext +from mstar.engine.resources.step import StepContext from mstar.model.nemotron_duplex.components.conformer import encoder_state_blocks, front_state_blocks from mstar.model.nemotron_duplex.config import ( ENCODER_FRONT, @@ -390,27 +389,43 @@ def test_codec_cleanup_clears_per_request_context(): assert "r" not in codec.request_states -class _EagerDecodeRegion: - """The engine's piecewise runner on CPU for the codec's decode region: the - batch padded to its bucket with zero codes, the captured callable run eagerly.""" +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 - def __init__(self, codec, bs: int = 4): - self.bs, self.calls = bs, 0 - self.config = codec.get_piecewise_accelerator_graph_configs(torch.device("cpu"), torch.float32)[ - AudioCodecDecoderSubmodule.DECODE_REGION] + from mstar.engine import accelerator_graph_runner as cgr + from mstar.engine.accelerator_graph_backend import AcceleratorGraphPool - def can_run(self, batch_size): - return batch_size <= self.bs + class _Replay: + def __init__(self, fn): + self.fn = fn - def run(self, static_inputs, real_bs): - self.calls += 1 - static = self.config.make_static_inputs(PiecewiseCaptureShape(bs=self.bs, seq_lens=[1] * self.bs, - total_tokens=self.bs)) - for name, value in static_inputs.items(): - static[name][:real_bs].copy_(value) - out = self.config.capture_fn(PiecewiseCallInputs( - static_inputs=static, engine_inputs=ModelInputsFromEngine(request_ids=[], per_request_info={}))) - return {name: value[:real_bs].clone() for name, value in out.items()} + 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), + ) + with mock.patch.object(cgr, "capture_into_graph", lambda fn, *a: (_Replay(fn), fn())): + runner.prepare_for_capture() + runner.warmup_and_capture() + 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]) @@ -436,7 +451,9 @@ def test_codec_stream_is_the_full_history_decode_with_the_real_decoder(frames, r runners = {} if region: cfg.eartts.num_quantizers = ccfg.num_quantizers - runners[AudioCodecDecoderSubmodule.DECODE_REGION] = _EagerDecodeRegion(codec) + 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(): @@ -592,44 +609,10 @@ def offline(self, wav): return self.node.perception(wav.unsqueeze(0), torch.tensor([wav.numel()]))[0][0] -class _EagerRegion: - """The engine's piecewise runner on CPU: the same bucket padding, region - declaration and admit / plan / commit around the region, with the captured - callable run eagerly.""" - - def __init__(self, enc: _Encoder, bs: int, label: str): - self.enc, self.bs, self.calls = enc, bs, 0 - self.config = enc.node.get_piecewise_accelerator_graph_configs(torch.device("cpu"), torch.float32)[label] - self.dummies = [f"__pad_{label}_{i}__" for i in range(bs)] - for rid in self.dummies: - for pool in enc.resources.values(): - pool.ingest_request(rid) - - def can_run(self, batch_size): - return batch_size <= self.bs - - def run(self, static_inputs, request_ids): - self.calls += 1 - n = len(request_ids) - static = self.config.make_static_inputs(PiecewiseCaptureShape(bs=self.bs, seq_lens=[1] * self.bs, - total_tokens=self.bs)) - for name, value in static_inputs.items(): - static[name][:n].copy_(value) - ids = [*request_ids, *self.dummies[n:]] - step = self.config.declare_step(ids, [1] * self.bs) - ctx = StepContext(request_ids=tuple(request_ids), graph_walk="pw", slot=0, capture=False, - slot_lease=SlotLease(slot=0, bucket=BucketKey("pw", self.bs, self.bs))) - ctx.set_padded_rids(tuple(ids)) - step.set_ctx(ctx) - assert self.enc.runner.admit(step).ok - self.enc.runner.plan(step) - out = self.config.capture_fn(PiecewiseCallInputs( - static_inputs=static, - engine_inputs=ModelInputsFromEngine(request_ids=ids, per_request_info={}, - resources=dict(self.enc.resources)), - )) - self.enc.runner.commit(step) - return {name: value[:n].clone() for name, value in out.items()} +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): @@ -706,9 +689,8 @@ def test_encoder_live_steps_replay_the_regions_with_padding_rows(): 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() - step = _EagerRegion(enc, bs=4, label=ConformerEncoderSubmodule.STEP_REGION) - layers = _EagerRegion(enc, bs=4, label=ConformerEncoderSubmodule.LAYERS_REGION) - enc.regions = {ConformerEncoderSubmodule.STEP_REGION: step, ConformerEncoderSubmodule.LAYERS_REGION: layers} + 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 From 1cd9bd620537e672f882747ac5b6353691cea0ff Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 04:05:44 -0700 Subject: [PATCH 08/69] R1.4: the encoder reads each pool's plan where it uses it, before a region re-plans The forward fetched both pools' plans up front, replayed the whole-step region for the calls one frame on (which re-plans both pools), and then ran the eager front end and layers on the plans it had fetched. That was correct only because the pool keeps eager and captured addressing in separate buffers; a forward holding a plan across another is what invariant 1 rules out. The eager pass now runs first and reads each plan when it uses it; the whole-step region runs last. A live step where every call is one frame on takes the region's lease (lease_before_step), and declare_step then declares nothing: the region's own step owns both pools, and no second plan is built and discarded. Every such call already holds its slots from an earlier step, so the region's admit, which cannot refuse, reserves nothing. A mixed step still declares every call, since a region's lease names the whole batch's bucket and its admit can't refuse a call without a slot; the comment in declare_step says so. Co-Authored-By: Claude Opus 5.5 --- mstar/model/nemotron_duplex/submodules.py | 105 ++++++++++++------ .../test_nemotron_duplex_submodules.py | 65 ++++++++++- 2 files changed, 127 insertions(+), 43 deletions(-) diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index e60b456292..c9147799b8 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -510,14 +510,29 @@ def _declare(request_ids, keys) -> SubmoduleStep: ) def declare_step(self, graph_walk, request_ids, inputs, slot_lease=None, piecewise_leases=None, - **kwargs) -> SubmoduleStep: - # Always the eager path's addressing. When the step replays a region - # instead, the region admits, plans and commits the same slots under - # its own step; for these pools that only re-reads the slot table and - # marks the slots written twice, so the two need no coordination. - del graph_walk, inputs, slot_lease, piecewise_leases, kwargs + **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 both pools. 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( + self._is_steady(rid, inp) for rid, inp in zip(request_ids, inputs, strict=True) + ): + return None + # Otherwise this step reserves and plans every call's slots, 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, and its admit + # can't refuse a call that has no slot yet. return self._declare(request_ids, (ENCODER_FRONT, ENCODER_STATE)) + def _is_steady(self, rid, inputs: NodeInputs) -> bool: + st = self.request_state(rid).get(self.STATE_KEY) + return st is not None and self.streamer.is_steady(st, inputs.tensor_inputs.get("chunk"), + inputs.kwargs["final"]) + @property def streamer(self) -> ConformerStreamer: if self._streamer is None: @@ -547,16 +562,16 @@ def layers_inputs(shape: PiecewiseCaptureShape) -> dict[str, torch.Tensor]: "n": torch.zeros(shape.bs, dtype=torch.long, device=device), } - def region(fn, inputs, keys) -> PiecewiseBatchedConfig: + 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), - seq_len=1, capture_batch_sizes=self.CAPTURE_BATCH_SIZES, compile=False, + seq_len=1, capture_batch_sizes=self.CAPTURE_BATCH_SIZES, compile=False, lease_before_step=lease, ) return { - self.STEP_REGION: region(self._step_region, step_inputs, (ENCODER_FRONT, ENCODER_STATE)), - self.LAYERS_REGION: region(self._layers_region, layers_inputs, (ENCODER_STATE,)), + 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), } def _step_region(self, call: PiecewiseCallInputs) -> dict[str, torch.Tensor]: @@ -578,6 +593,12 @@ def _stream_state(self, rid) -> StreamState: state.add(self.STATE_KEY, st) return st + @staticmethod + def _pool(engine_inputs, key) -> tuple: + """A pool and the plan this step made for it, read when the eager path uses it.""" + pool = engine_inputs.resources[key] + return pool, pool.addressing("main") + @staticmethod def _counts(values, device) -> torch.Tensor: """Host counts on the device without a blocking copy, which would wait @@ -592,44 +613,41 @@ def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, **k states = [self._stream_state(rid) for rid in rids] chunks = [c.to(device, non_blocking=True) if c is not None else None for c in chunk] final = list(final) - longest = int(self.BATCHED_CHUNK_S * self.perception.config.sample_rate) - front, pool = engine_inputs.resources[ENCODER_FRONT], engine_inputs.resources[ENCODER_STATE] - # this step's own addressing: a region's plan replaces the pools' current one - front_rows, pool_rows = front.addressing("main"), pool.addressing("main") runners = engine_inputs.piecewise_runners or {} frames: list[torch.Tensor | None] = [None] * len(rids) - # calls one frame of audio further on the frame grid: the whole step as one graph + # calls one frame of audio further on the frame grid: the whole step as one graph, run last step_runner = runners.get(self.STEP_REGION) steady = [i for i, (st, c, f) in enumerate(zip(states, chunks, final, strict=True)) if streamer.is_steady(st, c, f)] if step_runner is not None else [] - if steady and step_runner.can_run(len(steady)): - y = step_runner.run( - static_inputs={"chunk": torch.stack([chunks[i] for i in steady]), - "frames": self._counts([states[i].frames for i in steady], device)}, - 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] - chunks[i] = None # consumed; the rest of the step leaves the call as it is - - # A pass pads every call to its longest chunk, so a chunk of more than a - # second (an upload, a large append) runs alone, on its own row of the - # plan; the batched pass then leaves that call as it is. + if engine_inputs.step is None: + # declare_step left both pools to the region (see there) + assert len(steady) == len(rids), "an encoder step declared nothing but is not all steady" + elif steady and not step_runner.can_run(len(steady)): + steady = [] + steady_chunks = [chunks[i] for i in steady] + for i in steady: + chunks[i] = None # 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. A pass pads every call to + # its longest chunk, so a chunk of more than a second (an upload, a large + # append) runs alone, on its own row of the plan, and the batched pass + # leaves that call as it is. + longest = int(self.BATCHED_CHUNK_S * self.perception.config.sample_rate) for i, c in enumerate(chunks): if c is not None and c.numel() > longest: + front, front_rows = self._pool(engine_inputs, ENCODER_FRONT) enc = streamer.front([states[i]], [c], [final[i]], front, front_rows.narrow(i)) if enc is not None: - frames[i] = self._eager_layers(enc, pool, pool_rows.narrow(i), device)[0] + pool, rows = self._pool(engine_inputs, ENCODER_STATE) + frames[i] = self._eager_layers(enc, pool, rows.narrow(i), device)[0] streamer.advance([states[i]], enc) chunks[i], final[i] = None, False - - # the rest: the same math eagerly, the layers as a graph when no call got more than a frame - enc = streamer.front(states, chunks, final, front, front_rows) + # the batched pass: its layers as a graph when no call got more than a frame + enc = None + if len(steady) < len(rids): + enc = streamer.front(states, chunks, final, *self._pool(engine_inputs, ENCODER_FRONT)) if enc is not None: runner = runners.get(self.LAYERS_REGION) if runner is not None and enc.x.shape[1] == 1 and runner.can_run(len(rids)): @@ -638,11 +656,24 @@ def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, **k request_ids=rids)["y"] y = [row[:n] for row, n in zip(y, enc.n, strict=True)] else: - y = self._eager_layers(enc, pool, pool_rows, device) + y = self._eager_layers(enc, *self._pool(engine_inputs, ENCODER_STATE), device) streamer.advance(states, enc) for i, n in enumerate(enc.n): if n: frames[i] = y[i] + + if steady: + y = step_runner.run( + static_inputs={"chunk": torch.stack(steady_chunks), + "frames": self._counts([states[i].frames for i in steady], device)}, + 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. diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index cff1a8a57b..cfa9bb950c 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -582,22 +582,30 @@ def prepare(self, handle, chunk, final=False): 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.""" + leases = {label: lease for label, region in self.regions.items() + if region.lease_before_step and (lease := region.lease_slot(len(inputs))) is not None} + return self.node.declare_step("encode", rids, inputs, piecewise_leases=leases) + def step(self, rids, inputs) -> BatchedModelOutput: for rid in rids: self.pool.ingest_request(rid) self.front.ingest_request(rid) - step = self.node.declare_step("encode", rids, inputs) - 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) + 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), + 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)) - self.runner.commit(step) + if step is not None: + self.runner.commit(step) return out def encode(self, handle, chunk, final=False): @@ -698,6 +706,51 @@ def test_encoder_live_steps_replay_the_regions_with_padding_rows(): 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): + """Every read and write of a pool's state, eager or in a region, uses the + plan the pool holds at that moment, never one held across another plan: + the eager path runs on this step's plan before a region plans its own. A + call with a chunk over a second runs on a narrow of that plan.""" + enc = _Encoder() + enc.regions = _encoder_regions(enc, ConformerEncoderSubmodule.STEP_REGION, ConformerEncoderSubmodule.LAYERS_REGION) + for pool in enc.resources.values(): + for name in ("gather", "scatter_"): + def checked(block, rows, *args, _pool=pool, _op=getattr(pool, name), **kwargs): + plan = _pool.addressing("main").slot_indices + assert rows.slot_indices.untyped_storage().data_ptr() == plan.untyped_storage().data_ptr(), ( + "a pool read through a plan it no longer holds") + return _op(block, rows, *args, **kwargs) + + monkeypatch.setattr(pool, name, 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 test_encoder_an_all_steady_step_leaves_both_pools_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 pools: the node declares + nothing, so no second plan is built and thrown away. Otherwise it declares + every call in both pools.""" + 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)] + assert set(enc.declare([1, 2], nxt).keys()) == {ENCODER_FRONT, ENCODER_STATE} # no region captured + 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)] + assert set(enc.declare([1, 2], mixed).keys()) == {ENCODER_FRONT, ENCODER_STATE} + + @pytest.mark.parametrize("b_end_rides", [False, True]) @pytest.mark.parametrize("a_chunk", [1280, 9000]) @pytest.mark.parametrize("b_len", [16000 * 2, 16000 * 2 + 640, 16000 * 2 + 777]) From 49d6cc7a73a888c2b81e3c9c8337ae7b0d2352a9 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 04:06:58 -0700 Subject: [PATCH 09/69] R1.5: store the encoder's layer state in the dtype the layers compute it in ENCODER_STATE was declared bf16 whatever the model computed in. Under fp32 compute each step rounded its keys / values / conv inputs to bf16 for the next, so live frames drifted from one-shot ones (5e-4 relative instead of 2.5e-7), against what docs/models.rst says of fp32 numerics. The pool now takes the model's autocast dtype, and a deployment that sets another autocast_dtype is refused at startup, naming the pool: the pool is declared before the YAML is read, so it cannot follow an override. Co-Authored-By: Claude Opus 5.5 --- .../nemotron_duplex/nemotron_duplex_model.py | 13 ++++++++++--- test/modular/test_nemotron_duplex_model.py | 15 +++++++++++++++ 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/mstar/model/nemotron_duplex/nemotron_duplex_model.py b/mstar/model/nemotron_duplex/nemotron_duplex_model.py index cb2334f0c9..43ccce7889 100644 --- a/mstar/model/nemotron_duplex/nemotron_duplex_model.py +++ b/mstar/model/nemotron_duplex/nemotron_duplex_model.py @@ -457,7 +457,8 @@ def talker_kv_pages(self, page_size: int, sessions: int | None = None) -> int: return sessions * 2 * -(-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. + """Refuse an autocast dtype the encoder's state is not stored in, and 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``) and its talker pages from its first @@ -465,6 +466,12 @@ def validate_config_yaml(self, config: dict, config_path: str) -> None: another ends, and its frames with it. So ``max_concurrent_requests`` sessions must fit each; the conductor queues the rest. """ + dtype, state_dtype = config.get("autocast_dtype"), str(self.get_autocast_dtype()).removeprefix("torch.") + if dtype is not None and str(dtype).removeprefix("torch.") != state_dtype: + raise ValueError( + f"{config_path}: autocast_dtype {dtype} is not the {state_dtype} that " + f"resources.{ENCODER_STATE} stores the encoder's layer state in" + ) cap = config.get("max_concurrent_requests") if cap is None: logger.warning( @@ -517,8 +524,8 @@ def get_node_resources(self) -> list[NodeResourceSpec]: resource_key=ENCODER_STATE, nodes={"conformer_encoder"}, config=RecurrentStateConfig( num_layers=self.config.stt.num_layers, - # bf16, the engine's autocast dtype, which the layers compute the state in - blocks=encoder_state_blocks(self.config.stt, torch.bfloat16), + # in the dtype the layers compute it in, so a step stores what it computed + blocks=encoder_state_blocks(self.config.stt, self.get_autocast_dtype()), max_slots=self.DEFAULT_MAMBA_SLOTS + 1, ), ), diff --git a/test/modular/test_nemotron_duplex_model.py b/test/modular/test_nemotron_duplex_model.py index f45644c8b4..75470140f3 100644 --- a/test/modular/test_nemotron_duplex_model.py +++ b/test/modular/test_nemotron_duplex_model.py @@ -157,6 +157,21 @@ def test_duplex_refuses_a_session_cap_one_per_session_resource_cannot_hold(key): model.validate_config_yaml({"max_concurrent_requests": 48, "resources": resources}, "c.yaml") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_duplex_encoder_state_is_stored_in_the_compute_dtype(dtype, monkeypatch): + """The encoder stores each step's layer state for the next; in another + dtype than the layers compute it in, live frames drift from one-shot ones. + A deployment that computes in another dtype is refused at startup.""" + model = _make_model() + monkeypatch.setattr(model, "get_autocast_dtype", lambda: dtype) + (spec,) = [s for s in model.get_node_resources() if s.resource_key == ENCODER_STATE] + assert {block.dtype for block in spec.config.blocks.values()} == {dtype} + model.validate_config_yaml({"autocast_dtype": str(dtype).removeprefix("torch.")}, "c.yaml") + other = "float16" + with pytest.raises(ValueError, match=f"resources.{ENCODER_STATE}"): + model.validate_config_yaml({"autocast_dtype": other}, "c.yaml") + + def test_duplex_request_resource_configs(): model = _make_model() cfg = model.get_request_resource_configs({}, None) From 01a07ef1ec2a73619a6c796538394c338ce550aa Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 06:50:58 -0700 Subject: [PATCH 10/69] R1.6: the layer region runs only the calls that got a frame In a mixed step (calls one frame on in the whole-step region, the others in the eager front end), the layer region then ran every call of the step, those without a new frame included, so the 24-layer stack ran twice at the full bucket. It now runs only the calls with a frame, planning their rows itself; with one off-grid client among 31 on the grid that is a bucket of one, not of 32. The eager layer pass (a step where some call gets more than a frame) still pads every call to the widest; since a chunk of more than a second runs on its own row of the plan, that is at most 14 frames, not a whole upload's. Co-Authored-By: Claude Opus 5.5 --- mstar/model/nemotron_duplex/submodules.py | 18 ++++++++++-------- .../modular/test_nemotron_duplex_submodules.py | 17 +++++++++++++++++ 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index c9147799b8..b8a173140e 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -650,17 +650,19 @@ def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, **k enc = streamer.front(states, chunks, final, *self._pool(engine_inputs, ENCODER_FRONT)) if enc is not None: runner = runners.get(self.LAYERS_REGION) - if runner is not None and enc.x.shape[1] == 1 and runner.can_run(len(rids)): - counts = self._counts([enc.frames, enc.n], device) - y = runner.run(static_inputs={"x": enc.x, "frames": counts[0], "n": counts[1]}, - request_ids=rids)["y"] - y = [row[:n] for row, n in zip(y, enc.n, strict=True)] + new = [i for i, n in enumerate(enc.n) if n] + if runner is not None and enc.x.shape[1] == 1 and runner.can_run(len(new)): + # only the calls with a frame: the region plans its own rows + counts = self._counts([[enc.frames[i] for i in new], [1] * len(new)], device) + y = runner.run(static_inputs={"x": torch.stack([enc.x[i] for i in new]), + "frames": counts[0], "n": counts[1]}, + request_ids=[rids[i] for i in new])["y"] + y = dict(zip(new, y, strict=True)) else: y = self._eager_layers(enc, *self._pool(engine_inputs, ENCODER_STATE), device) streamer.advance(states, enc) - for i, n in enumerate(enc.n): - if n: - frames[i] = y[i] + for i in new: + frames[i] = y[i][: enc.n[i]] if steady: y = step_runner.run( diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index cfa9bb950c..91c5676281 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -706,6 +706,23 @@ def test_encoder_live_steps_replay_the_regions_with_padding_rows(): 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): """Every read and write of a pool's state, eager or in a region, uses the plan the pool holds at that moment, never one held across another plan: From 493448f977e34d4afa2b7d46811c06bc8b53f819 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 04:10:48 -0700 Subject: [PATCH 11/69] R1.10: test the recurrent pool's gather / scatter_, and refuse the sentinel gather / scatter_ came over from #383 without its test, and the num_rows comment still said it counts real rows only, while both methods index every declared row, a replay's padding included. On a pool built with disable_sink_slot those padding rows carry the -1 sentinel, which the methods then indexed with: an IndexError on CPU, a device-side assert on CUDA. Bring #383's test and comment fix (the methods' bodies stay as in #383), and refuse padded addressing with a plain error on a pool with no sink, checked on the host plan and only for such pools. docs/adding_models.rst gains a row for RecurrentStateSpec, gather and scatter_ included. Co-Authored-By: Claude Opus 5.5 --- docs/adding_models.rst | 11 +++++ mstar/engine/resources/recurrent/pool.py | 14 +++++- test/modular/test_recurrent_padding_rows.py | 53 +++++++++++++++++++++ 3 files changed, 77 insertions(+), 1 deletion(-) diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 927e0e277a..a11e6f9aa3 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -434,6 +434,17 @@ The spec types are: ``high_freq_factor`` and ``old_context_len``. A model with learned position embeddings also declares this spec, because it needs the position counter. Such a model never applies RoPE. + * - ``RecurrentStateSpec(config=RecurrentStateConfig(...))`` + - A pool of fixed-size per-request state slots (a Mamba or gated-DeltaNet + layer's conv window and SSM state, a streaming encoder's caches), one slot per + request and label, held until the request ends. ``blocks`` names each tensor of + a slot and its shape and dtype; ``max_slots`` (one more than the requests, for + the sink that padding rows address) is deployment-tunable. A recurrent kernel + reads the step's ``addressing``; plain torch reads and writes a block through + ``gather`` / ``scatter_`` with that addressing (or its ``narrow``, for rows a + forward runs apart from the rest), so eager and captured forwards + address the same slots. Those two need the sink: a pool built with + ``disable_sink_slot`` refuses them for a padded step. * - ``SamplerSpec(vocab_size=..., enable_repetion_penalty=...)`` - A sampler with its own per-request parameter buffers, philox stream, and optional seen-token mask. ``enable_repetion_penalty`` declares a capability, not an diff --git a/mstar/engine/resources/recurrent/pool.py b/mstar/engine/resources/recurrent/pool.py index 3f1e0a68ac..d09e173cf4 100644 --- a/mstar/engine/resources/recurrent/pool.py +++ b/mstar/engine/resources/recurrent/pool.py @@ -59,7 +59,7 @@ class RecurrentAddressing: slot_indices: torch.Tensor # [rows] int32, `pad_index` where nothing is addressed has_state: torch.Tensor # [rows] bool, False where the slot reads as zeros - num_rows: int # real rows, before a capture bucket's padding + num_rows: int # the rows the step declared, a replay's padding rows included # The same two on the host, as built: a backend that splits rows between # kernels by them must not read the device copies back host_slot_indices: Sequence[int] | None = None @@ -158,6 +158,7 @@ def gather(self, name: str, rows: RecurrentAddressing, layer_idx: int = 0) -> to ``rows`` is the step's addressing (``addressing``), so a captured and an eager forward address their slots the same way; padding rows read the sink.""" + self._check_addressable(rows) index = rows.slot_indices[: rows.num_rows].long() return self._blocks[name][layer_idx].index_select(0, index) @@ -165,9 +166,20 @@ def scatter_( self, name: str, rows: RecurrentAddressing, value: torch.Tensor, layer_idx: int = 0, ) -> None: """Write ``value [rows, *shape]`` into block ``name`` of these rows' slots.""" + self._check_addressable(rows) index = rows.slot_indices[: rows.num_rows].long() self._blocks[name][layer_idx].index_copy_(0, index, value) + def _check_addressable(self, rows: RecurrentAddressing) -> None: + # Without a sink, padding rows carry the sentinel, which only kernels + # that skip rows honour; indexing with it fails on the device. Checked + # on the host plan, and only for pools without a sink. + if self.config.disable_sink_slot and NO_SLOT in (rows.host_slot_indices or ()): + raise ValueError( + "gather / scatter_ address padding rows at the sink, and this pool " + "has none (disable_sink_slot)" + ) + @property def num_free_slots(self) -> int: return len(self._free) diff --git a/test/modular/test_recurrent_padding_rows.py b/test/modular/test_recurrent_padding_rows.py index ebbbc43a24..02a3b2ce7b 100644 --- a/test/modular/test_recurrent_padding_rows.py +++ b/test/modular/test_recurrent_padding_rows.py @@ -13,6 +13,7 @@ from __future__ import annotations +import pytest import torch from mstar.engine.resources.base import EngineResourceInfo @@ -147,6 +148,58 @@ def test_a_real_request_padding_and_release_across_the_rotation(): assert pool.num_free_slots == 2 +def test_gather_scatter_follow_the_step_addressing(): + """``gather`` / ``scatter_`` address the step's rows: a replay's padding rows + read and write the sink and nothing else, and a ``select`` of the step + addresses the same slots as the step.""" + pool = build_pool(max_slots=6) + step = decode_step(["a", "b"]) + eager = StepContext(request_ids=["a", "b"], graph_walk="decode", slot=0, capture=False) + pool.admit(step, eager) + rows = pool.plan(step, eager)["main"] + value = torch.arange(2 * 96 * 3, dtype=torch.float32).view(2, 96, 3) + pool.scatter_("conv", rows, value, layer_idx=1) + torch.testing.assert_close(pool.gather("conv", rows, layer_idx=1), value) + one = StepContext(request_ids=["b"], graph_walk="decode", slot=0, capture=False) + torch.testing.assert_close(pool.gather("conv", pool.plan(decode_step(["b"]), one)["main"], layer_idx=1), value[1:]) + untouched = pool.block("conv", 0).clone() + + padded = decode_step(["a", "b", "pad0", "pad1"]) + ctx = padded_ctx(["a", "b"], ["a", "b", "pad0", "pad1"], bs=4) + pool.admit(padded, ctx) + rows = pool.plan(padded, ctx)["main"] + assert rows.num_rows == 4 + got = pool.gather("conv", rows, layer_idx=1) + torch.testing.assert_close(got[:2], value) + pool.scatter_("conv", rows, torch.full((4, 96, 3), -1.0), layer_idx=1) + conv = pool.block("conv", 1) + live = {pool._slots[r]["main"].index for r in ("a", "b")} + assert all(bool((conv[i] == -1).all()) for i in live | {SINK_SLOT}) + assert all(bool((conv[i] == 0).all()) for i in range(6) if i not in live | {SINK_SLOT}) + torch.testing.assert_close(pool.block("conv", 0), untouched) + + +def test_gather_scatter_refuse_padding_rows_on_a_pool_without_a_sink(): + """Padding rows of a pool with no sink carry the sentinel, which no slot + holds: ``gather`` / ``scatter_`` say so rather than index with it (an + IndexError on CPU, a device-side assert on CUDA).""" + pool = build_pool(max_slots=4, disable_sink_slot=True) + names = [f"__cg_LLM_0_{i}__" for i in range(4)] + for rid in ["r", *names]: + pool.ingest_request(rid) + padded = ["r", *names[1:]] + ctx = padded_ctx(["r"], padded, bs=4) + assert pool.admit(decode_step(padded), ctx).ok + rows = pool.plan(decode_step(padded), ctx)["main"] + with pytest.raises(ValueError, match="sink"): + pool.gather("conv", rows) + with pytest.raises(ValueError, match="sink"): + pool.scatter_("conv", rows, torch.zeros(4, 96, 3)) + eager = StepContext(request_ids=["r"], graph_walk="decode", slot=0, capture=False) + real = pool.plan(decode_step(["r"]), eager)["main"] + assert pool.gather("conv", real).shape == (1, 96, 3) + + def test_narrow_addresses_some_rows_of_the_step(): """``narrow`` is a run of the step's rows, as views of its plan: a forward that runs one row apart reads and writes that row's slot and no other.""" From e6e08823c23792c321830f1dc5fe94755f2a5548 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 04:11:44 -0700 Subject: [PATCH 12/69] R1.11: gather reads with the plan's int32 indices, without a cast The encoder's layer loop gathers three blocks a layer, 24 layers a step, and each gather first cast the plan's int32 slot indices to int64: a kernel apiece that index_select does not need. index_copy_ does need int64, so scatter_ keeps its cast. Not measured on a GPU; it is one kernel fewer per gather, 72 a layer pass. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/recurrent/pool.py | 4 ++-- test/modular/test_recurrent_padding_rows.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/mstar/engine/resources/recurrent/pool.py b/mstar/engine/resources/recurrent/pool.py index d09e173cf4..5ccca0963d 100644 --- a/mstar/engine/resources/recurrent/pool.py +++ b/mstar/engine/resources/recurrent/pool.py @@ -159,8 +159,8 @@ def gather(self, name: str, rows: RecurrentAddressing, layer_idx: int = 0) -> to eager forward address their slots the same way; padding rows read the sink.""" self._check_addressable(rows) - index = rows.slot_indices[: rows.num_rows].long() - return self._blocks[name][layer_idx].index_select(0, index) + # index_select takes the int32 indices as they are; index_copy_ needs int64 + return self._blocks[name][layer_idx].index_select(0, rows.slot_indices[: rows.num_rows]) def scatter_( self, name: str, rows: RecurrentAddressing, value: torch.Tensor, layer_idx: int = 0, diff --git a/test/modular/test_recurrent_padding_rows.py b/test/modular/test_recurrent_padding_rows.py index 02a3b2ce7b..391f2925ee 100644 --- a/test/modular/test_recurrent_padding_rows.py +++ b/test/modular/test_recurrent_padding_rows.py @@ -200,6 +200,19 @@ def test_gather_scatter_refuse_padding_rows_on_a_pool_without_a_sink(): assert pool.gather("conv", real).shape == (1, 96, 3) +def test_gather_reads_with_the_planned_indices_as_they_are(): + """A layer loop gathers several blocks a layer, every step: each read is the + index_select alone, with no cast of the int32 plan to int64 first.""" + pool = build_pool(max_slots=4) + ctx = StepContext(request_ids=["a"], graph_walk="decode", slot=0, capture=False) + pool.admit(decode_step(["a"]), ctx) + rows = pool.plan(decode_step(["a"]), ctx)["main"] + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU], acc_events=True) as prof: + pool.gather("conv", rows) + ops = {event.name for event in prof.events()} + assert "aten::index_select" in ops and not ops & {"aten::to", "aten::_to_copy"}, ops + + def test_narrow_addresses_some_rows_of_the_step(): """``narrow`` is a run of the step's rows, as views of its plan: a forward that runs one row apart reads and writes that row's slot and no other.""" From 12ddec7d12d8243eb64ace3cc5f0d63134f73763 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 04:12:44 -0700 Subject: [PATCH 13/69] R1.7: describe the encoder's pools and graphs as they are The module and class docstrings still said the encoder declares no resources and runs eager, that a call keeps all its mel frames in its per-request state, and that the layer stack is the only captured region; the model's session comment named only mamba_state; docs/models.rst described the first commit's layer pool alone and gave the live-lag numbers from before the captured encoder as current. Each now says what the code does, with the PR's measured numbers for the captured encoder. Co-Authored-By: Claude Opus 5.5 --- docs/models.rst | 24 +++++++++++----- .../nemotron_duplex/nemotron_duplex_model.py | 13 +++++---- mstar/model/nemotron_duplex/submodules.py | 28 ++++++++++++------- 3 files changed, 43 insertions(+), 22 deletions(-) diff --git a/docs/models.rst b/docs/models.rst index ccab3c2208..205a299919 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -429,7 +429,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 @@ -472,8 +473,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. @@ -520,10 +525,15 @@ 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 layer state (its attention keys / values and - conv inputs) in a recurrent-state pool, a slot per session; a live step, one - frame per call, replays its layer stack as a CUDA graph, and a step with more - frames (a one-shot upload, a large append) runs it eagerly. +- The encoder keeps each call's state in two recurrent-state pools, a slot per + session each: its layer state (attention keys / values and conv inputs, + ``encoder_state``) and its newest audio and mel frames (``encoder_front``). 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; other calls + take the front end eagerly, then the layer stack as a graph when each got at + most one frame, or eagerly (a one-shot upload, a large append, the call's + end). 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/mstar/model/nemotron_duplex/nemotron_duplex_model.py b/mstar/model/nemotron_duplex/nemotron_duplex_model.py index 43ccce7889..c186e79bf4 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 two recurrent-state pools (each call's layer state, 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, @@ -412,8 +414,9 @@ def _talker_tokenizer(self): # 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``. + # Sessions resident at once (a slot each in mamba_state, ~137 MB in fp32 for + # Nemotron-H, and in encoder_state / encoder_front; pages in talker_kv). A + # deployment raising it raises all four (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_MAMBA_SLOTS = 32 # sessions; the runner's padding rows take no slot (#258) diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index b8a173140e..67888e392c 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -2,7 +2,8 @@ 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_state, 16 kHz speech -> per-frame LLM embeds (+ RNN-T); per + encoder_front session, its layers' state 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 +12,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 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 @@ -433,19 +436,24 @@ class ConformerEncoderSubmodule(NodeSubmodule): 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. - - A call's layer state (attention keys / values, conv inputs) lives in the - ``ENCODER_STATE`` recurrent pool; its mel front end in its per-request - state. A live step gives every call at most one frame, and its layer stack - replays a CUDA graph (``LAYERS_REGION``); a step with more (a one-shot - request's whole clip, a large append, the call's end) runs it eagerly. + pass of its own within the same engine step, on its own row of the step's + plan. A call's state is bounded whatever its length: its newest audio and + mel frames (the ``ENCODER_FRONT`` pool), and the last 70 frames of each + layer's attention keys and values plus the conv tails (``ENCODER_STATE``); + its host counters live in its per-request state. + + 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, else eagerly (a + one-shot upload, a large append, the call's end). Holds the perception stack plus the RNN-T decoder/joint (loaded; the user-transcript channel is not emitted yet). """ - # Variable-length front end, eager; the layer stack is the captured region. + # The captured regions are not compiled (compile=False); the eager front end + # takes variable-length chunks. disable_torch_compile = True STEP_REGION = "conformer_step" From a115c768254807d31688b74b0d354938affcc66e Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 04:13:15 -0700 Subject: [PATCH 14/69] R1.13: tell live clients which append size takes the encoder's fast path The whole-step graph serves only a call appending exactly one 80 ms frame (1280 samples) on its frame grid; a client appending 20, 40 or 100 ms is served the same, through the eager front end. docs/clients.rst now says so next to the live-input protocol (docs/models.rst says it in the R1.7 commit). Generalizing the fast path to other sizes is left out of this PR. Co-Authored-By: Claude Opus 5.5 --- docs/clients.rst | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/clients.rst b/docs/clients.rst index 713a46c3b6..e37ebf0cb3 100644 --- a/docs/clients.rst +++ b/docs/clients.rst @@ -177,6 +177,9 @@ not the end of a request, and the call goes on. A call still waiting for a slot plus one each for the opening message's media and the end of input; the next ends the call with ``status`` 503. 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 a call whose appends are exactly 80 ms (1280 samples at 16 kHz) on its +frame grid; other sizes are served the same, more slowly (see :doc:`models`). .. code-block:: python From 76f6b7f322b4a2e5ebdef5dde9b8951cc4150387 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 04:13:15 -0700 Subject: [PATCH 15/69] R1.15: say that per-message deflate is off for every WebSocket route The change turning deflate off said it was for /generate/ws, but uvicorn negotiates it per server, so /v1/realtime lost it too. uvicorn has no per-route setting, and both routes carry audio that barely compresses, so the setting stays server-wide; the comment and docs/clients.rst now say so. Co-Authored-By: Claude Opus 5.5 --- docs/clients.rst | 2 ++ mstar/api_server/entrypoint.py | 6 ++++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/docs/clients.rst b/docs/clients.rst index e37ebf0cb3..e553aeae89 100644 --- a/docs/clients.rst +++ b/docs/clients.rst @@ -86,6 +86,8 @@ 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 no per-message deflate on any WebSocket route: the audio they mostly +carry barely compresses, and compressing it costs the event loop. 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 diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index 700d3a9f26..39d679e045 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -2481,8 +2481,10 @@ def main(argv: list[str] | None = None): server = uvicorn.Server( uvicorn.Config( app, host=args.host, port=args.port, access_log=False, - # /generate/ws streams PCM, which deflate barely shrinks; - # compressing it cost the event loop ~8% at 32 live calls + # Server-wide (uvicorn negotiates it for every WebSocket + # route): /generate/ws and /v1/realtime stream audio, which + # deflate barely shrinks, and compressing it cost the event + # loop ~8% at 32 live calls on /generate/ws ws_per_message_deflate=False, ) ) From 5ee201ad36a2d7f6512b83ce80c95ad8ffa99d0b Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 04:15:13 -0700 Subject: [PATCH 16/69] R1.17: one list of session batch sizes; the codec reads its runners directly Four nodes each spelled out [1, 2, 4, 8, 16, 32] for their CUDA-graph batches, and DEFAULT_MAMBA_SLOTS said 32 on its own; they are now one SESSION_BATCH_SIZES in config.py, whose largest is the default session count. The codec's forward read piecewise_runners through getattr, for test fakes that lacked the field; the fakes now carry it. Co-Authored-By: Claude Opus 5.5 --- mstar/model/nemotron_duplex/config.py | 4 ++++ mstar/model/nemotron_duplex/nemotron_duplex_model.py | 3 ++- mstar/model/nemotron_duplex/submodules.py | 11 ++++++----- test/modular/test_nemotron_duplex_submodules.py | 6 +++--- 4 files changed, 15 insertions(+), 9 deletions(-) diff --git a/mstar/model/nemotron_duplex/config.py b/mstar/model/nemotron_duplex/config.py index 8abaa61966..5c88760096 100644 --- a/mstar/model/nemotron_duplex/config.py +++ b/mstar/model/nemotron_duplex/config.py @@ -33,6 +33,10 @@ ENCODER_STATE = "encoder_state" # recurrent-state pool: each call's streaming-encoder layer state ENCODER_FRONT = "encoder_front" # recurrent-state pool: each call's streaming-encoder audio and mel tails +# Every node's CUDA-graph batch sizes: a row per live session, up to the sessions +# a deployment serves by default (``NemotronDuplexModel.DEFAULT_MAMBA_SLOTS``). +SESSION_BATCH_SIZES = [1, 2, 4, 8, 16, 32] + def session_too_long(max_session_s: float) -> ValueError: """The one error for a call whose audio outgrows ``max_session_s``, whichever diff --git a/mstar/model/nemotron_duplex/nemotron_duplex_model.py b/mstar/model/nemotron_duplex/nemotron_duplex_model.py index c186e79bf4..a8270b5098 100644 --- a/mstar/model/nemotron_duplex/nemotron_duplex_model.py +++ b/mstar/model/nemotron_duplex/nemotron_duplex_model.py @@ -82,6 +82,7 @@ NANO_ATTN, NANO_KV, NANO_SAMPLER, + SESSION_BATCH_SIZES, TALKER_ATTN, TALKER_KV, TALKER_POS, @@ -419,7 +420,7 @@ def _talker_tokenizer(self): # deployment raising it raises all four (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_MAMBA_SLOTS = 32 # sessions; the runner's padding rows take no slot (#258) + DEFAULT_MAMBA_SLOTS = SESSION_BATCH_SIZES[-1] # 32 sessions; 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 # once: a full pool stalls every session, not just the one that outgrew it. diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index 67888e392c..1c6f7afbc2 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -57,6 +57,7 @@ NANO_ATTN, NANO_KV, NANO_SAMPLER, + SESSION_BATCH_SIZES, TALKER_ATTN, TALKER_CFG_LABEL, TALKER_KV, @@ -107,7 +108,7 @@ class NemotronHLLMSubmodule(ARNodeSubmodule): # 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] + 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 @@ -458,7 +459,7 @@ class ConformerEncoderSubmodule(NodeSubmodule): STEP_REGION = "conformer_step" LAYERS_REGION = "conformer_layers" - CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32] + 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 @@ -728,7 +729,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): @@ -963,7 +964,7 @@ class AudioCodecDecoderSubmodule(NodeSubmodule): disable_autocast = True DECODE_REGION = "codec_decode" - CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32] + CAPTURE_BATCH_SIZES = SESSION_BATCH_SIZES CONTEXT_KEY = "codec_ctx" TAIL_KEY = "codec_tail" # the last chunk's held-back samples, int16 @@ -1060,7 +1061,7 @@ def forward_batched( by_len: dict[int, list[int]] = {} for i, (full, _) in windows.items(): by_len.setdefault(full.shape[0], []).append(i) - runner = (getattr(engine_inputs, "piecewise_runners", None) or {}).get(self.DECODE_REGION) + 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) if runner is not None and tf == self._steady_window and runner.can_run(len(idxs)): diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index 91c5676281..4844ef04e2 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -300,12 +300,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) @@ -333,7 +333,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)) From 739b56d87d94858c727b6e38b75bb4019c9e85e2 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 07:03:07 -0700 Subject: [PATCH 17/69] R2.1: check the codec's window envelope on the host, so it builds on meta R1.1 put check_window_envelope in AudioCodec.__init__, which the model only calls under torch.device("meta"): the window was built on meta and the check's .all() had no values to read, so the codec node could not be built and the worker died at startup. The check's window is built on the CPU. A test builds every node's modules through the model's own get_submodule, under meta, as the worker does (the codec failed it before; the other three nodes passed). Co-Authored-By: Claude Opus 5.5 --- .../nemotron_duplex/components/audio_codec.py | 3 +- test/modular/test_nemotron_duplex_model.py | 30 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/mstar/model/nemotron_duplex/components/audio_codec.py b/mstar/model/nemotron_duplex/components/audio_codec.py index b6f6848454..9c006e3d0b 100644 --- a/mstar/model/nemotron_duplex/components/audio_codec.py +++ b/mstar/model/nemotron_duplex/components/audio_codec.py @@ -225,7 +225,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 - check_window_envelope(torch.hann_window(cfg.n_fft), cfg.hop_length) + # 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: diff --git a/test/modular/test_nemotron_duplex_model.py b/test/modular/test_nemotron_duplex_model.py index 75470140f3..1fc18f50a2 100644 --- a/test/modular/test_nemotron_duplex_model.py +++ b/test/modular/test_nemotron_duplex_model.py @@ -172,6 +172,36 @@ def test_duplex_encoder_state_is_stored_in_the_compute_dtype(dtype, monkeypatch) model.validate_config_yaml({"autocast_dtype": other}, "c.yaml") +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_request_resource_configs(): model = _make_model() cfg = model.get_request_resource_configs({}, None) From 44c8105d0670b5d39d211a55fd0a2eae9cffacb5 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 07:03:42 -0700 Subject: [PATCH 18/69] R2.2: cap the encoder's and codec's batches at their widest graph Engine.get_max_batch_size caps a node's batch at its whole-forward graphs only. The encoder and the codec capture piecewise regions alone, so their cap was None: past 32 sessions a step of 48 calls fit no bucket, and every call ran eagerly, every step. Both nodes now declare max_batch_size as their widest capture (32), so the scheduler splits a wider batch into steps that replay. Making the engine's cap see piecewise regions is the general fix; it is a system change, left as a follow-up. Co-Authored-By: Claude Opus 5.5 --- docs/models.rst | 3 +++ mstar/model/nemotron_duplex/submodules.py | 9 +++++++++ .../modular/test_nemotron_duplex_submodules.py | 18 ++++++++++++++++++ 3 files changed, 30 insertions(+) diff --git a/docs/models.rst b/docs/models.rst index 205a299919..bd8df13d7f 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -448,6 +448,9 @@ Nemotron VoiceChat (``nemotron_duplex``) notes to match; to serve more, raise the cap and all four 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 diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index 1c6f7afbc2..ac90647fa3 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -511,6 +511,11 @@ 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: + # 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) -> SubmoduleStep: return SubmoduleStep( @@ -1006,6 +1011,10 @@ 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 diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index 4844ef04e2..d88ce84aa3 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -819,6 +819,24 @@ def spy_layers(x, pool, rows, *rest): 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(), max_session_samples=16000), + "audio_codec": AudioCodecDecoderSubmodule(codec=nn.Identity(), config=NemotronDuplexConfig()), + } + engine = Engine.__new__(Engine) + engine._submodules = {name: SimpleNamespace(submodule=node, accelerator_graph_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 + + 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.""" From 4cfac81d15929979aa8a3872f2f089294ee21755 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 07:04:24 -0700 Subject: [PATCH 19/69] R2.3: size and check nano_kv for the session cap, as talker_kv is validate_config_yaml checked every per-session resource but nano_kv, which holds a session's prompt and every frame for life (no eviction by default): at cap 32, model_kwargs max_session_s: 640 needs 67 pages a session and the default 2048-page pool holds 30 sessions; max_prompt_tokens: 8192 holds 25; resources: {nano_kv: {max_num_pages: 2}} passed. A session past the pool waits on admit, and its frames with it. nano_kv is now sized like talker_kv, for 32 sessions of the longest prompt and call served (641 pages at the defaults, 20 a session, where it was a fixed 2048), and the startup check counts the sessions it holds. Both pools share one page count, prefix + max_session_frames + 1. Co-Authored-By: Claude Opus 5.5 --- docs/models.rst | 8 ++++--- .../nemotron_duplex/nemotron_duplex_model.py | 22 ++++++++++++----- test/modular/test_nemotron_duplex_model.py | 24 +++++++++++++++---- 3 files changed, 40 insertions(+), 14 deletions(-) diff --git a/docs/models.rst b/docs/models.rst index bd8df13d7f..1d06439345 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -443,9 +443,11 @@ Nemotron VoiceChat (``nemotron_duplex``) notes 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 its pages in ``talker_kv`` - (``max_num_pages``). The shipped configs set ``max_concurrent_requests: 32`` - to match; to serve more, raise the cap and all four with it. The server + more than the sessions, for the 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 five 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 diff --git a/mstar/model/nemotron_duplex/nemotron_duplex_model.py b/mstar/model/nemotron_duplex/nemotron_duplex_model.py index a8270b5098..f6d8150b3e 100644 --- a/mstar/model/nemotron_duplex/nemotron_duplex_model.py +++ b/mstar/model/nemotron_duplex/nemotron_duplex_model.py @@ -455,18 +455,25 @@ def max_session_frames(self) -> int: def talker_kv_pages(self, page_size: int, sessions: int | None = None) -> int: """Pages for ``sessions`` (DEFAULT_MAMBA_SLOTS) 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_MAMBA_SLOTS) 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_MAMBA_SLOTS if sessions is None else sessions - positions = self.TALKER_WARMUP_POSITIONS + self.max_session_frames + 1 + positions = prefix + self.max_session_frames + 1 # + 1 for the pool's sink page - return sessions * 2 * -(-positions // page_size) + 1 + return sessions * streams * -(-positions // page_size) + 1 def validate_config_yaml(self, config: dict, config_path: str) -> None: """Refuse an autocast dtype the encoder's state is not stored in, and 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``) and its talker pages from its first - step to its last, so a session past any of them waits on admit until + ``encoder_state``, ``encoder_front``) 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. """ @@ -490,9 +497,10 @@ def validate_config_yaml(self, config: dict, config_path: str) -> None: for spec in specs: if isinstance(spec, RecurrentStateSpec): held, knob = spec.config.usable_slots, "max_slots" - elif spec.resource_key == TALKER_KV: + elif spec.resource_key in (TALKER_KV, NANO_KV): kv = spec.config - per_session = self.talker_kv_pages(kv.page_size, sessions=1) - 1 + 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 @@ -522,6 +530,8 @@ 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( diff --git a/test/modular/test_nemotron_duplex_model.py b/test/modular/test_nemotron_duplex_model.py index 1fc18f50a2..17b41b56cc 100644 --- a/test/modular/test_nemotron_duplex_model.py +++ b/test/modular/test_nemotron_duplex_model.py @@ -141,22 +141,36 @@ def test_duplex_shipped_configs_hold_their_session_cap(): model.validate_config_yaml(yaml.safe_load(path.read_text()), str(path)) -@pytest.mark.parametrize("key", [MAMBA_STATE, ENCODER_STATE, ENCODER_FRONT, TALKER_KV]) +@pytest.mark.parametrize("key", [MAMBA_STATE, ENCODER_STATE, ENCODER_FRONT, 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 talker pages - from start to end; a session past any of them waits on admit until + """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() - pages = 48 * (model.talker_kv_pages(128) - 1) // model.DEFAULT_MAMBA_SLOTS + 1 resources = {MAMBA_STATE: {"max_slots": 49}, ENCODER_STATE: {"max_slots": 49}, - ENCODER_FRONT: {"max_slots": 49}, TALKER_KV: {"max_num_pages": pages}} + ENCODER_FRONT: {"max_slots": 49}, 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("dtype", [torch.bfloat16, torch.float32]) def test_duplex_encoder_state_is_stored_in_the_compute_dtype(dtype, monkeypatch): """The encoder stores each step's layer state for the next; in another From 1d042b658a591eeba146906b9a79bc2db331f6d4 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 07:05:15 -0700 Subject: [PATCH 20/69] R2.4: the encoder's layer state takes the engine's compute dtype R1.5 declared ENCODER_STATE in the model's own autocast dtype (bf16) and refused a YAML autocast_dtype other than it. autocast_dtype: null, the one override that runs fp32 with no autocast, passed the check (it skipped None) while the pool stayed bf16, so live frames drifted from one-shot ones (4.6e-4 relative vs 2.0e-7 with an fp32 pool), a regression from #594, whose per-request state took the runtime dtype. A RecurrentBlockConfig may now leave its dtype None, which the pool settles when built to the engine's compute dtype, the one its KV caches store (EngineResourceInfo.kv_dtype; torch's default when no autocast runs). The encoder's layer state declares None, so it follows any deployment dtype, and the startup refusal goes. The pool change is engine-level but opt-in: blocks with a dtype are untouched. Co-Authored-By: Claude Opus 5.5 --- docs/adding_models.rst | 3 ++- mstar/engine/resources/recurrent/config.py | 5 +++-- mstar/engine/resources/recurrent/pool.py | 3 +++ .../nemotron_duplex/components/conformer.py | 7 +++++-- .../nemotron_duplex/nemotron_duplex_model.py | 12 ++--------- test/modular/test_nemotron_duplex_model.py | 20 ++++++++++--------- 6 files changed, 26 insertions(+), 24 deletions(-) diff --git a/docs/adding_models.rst b/docs/adding_models.rst index a11e6f9aa3..ffa879d2dc 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -438,7 +438,8 @@ The spec types are: - A pool of fixed-size per-request state slots (a Mamba or gated-DeltaNet layer's conv window and SSM state, a streaming encoder's caches), one slot per request and label, held until the request ends. ``blocks`` names each tensor of - a slot and its shape and dtype; ``max_slots`` (one more than the requests, for + a slot and its shape and dtype (None: the engine's compute dtype, what its KV + caches store); ``max_slots`` (one more than the requests, for the sink that padding rows address) is deployment-tunable. A recurrent kernel reads the step's ``addressing``; plain torch reads and writes a block through ``gather`` / ``scatter_`` with that addressing (or its ``narrow``, for rows a diff --git a/mstar/engine/resources/recurrent/config.py b/mstar/engine/resources/recurrent/config.py index 097b0a4f51..58a911f38a 100644 --- a/mstar/engine/resources/recurrent/config.py +++ b/mstar/engine/resources/recurrent/config.py @@ -24,11 +24,12 @@ 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), settled when the pool is built. """ shape: tuple[int, ...] - dtype: torch.dtype + dtype: torch.dtype | None shard_dims: tuple[int, ...] = () def __post_init__(self): diff --git a/mstar/engine/resources/recurrent/pool.py b/mstar/engine/resources/recurrent/pool.py index 5ccca0963d..1078d88d8c 100644 --- a/mstar/engine/resources/recurrent/pool.py +++ b/mstar/engine/resources/recurrent/pool.py @@ -88,6 +88,9 @@ def build( config = spec.config if info.joint_comm_group is not None: config.shard(info.joint_comm_group.world_size) + for block in config.blocks.values(): + if block.dtype is None: + block.dtype = info.kv_dtype if info.kv_dtype is not None else torch.get_default_dtype() return cls(device=info.device, config=config) def __init__(self, device: torch.device, config: RecurrentStateConfig): diff --git a/mstar/model/nemotron_duplex/components/conformer.py b/mstar/model/nemotron_duplex/components/conformer.py index d37e690d3f..4d81706a34 100644 --- a/mstar/model/nemotron_duplex/components/conformer.py +++ b/mstar/model/nemotron_duplex/components/conformer.py @@ -357,9 +357,12 @@ def front_state_blocks(config: ConformerSTTConfig) -> dict[str, RecurrentBlockCo } -def encoder_state_blocks(config: ConformerSTTConfig, dtype: torch.dtype) -> dict[str, RecurrentBlockConfig]: +def encoder_state_blocks(config: ConformerSTTConfig, + dtype: torch.dtype | None = None) -> dict[str, RecurrentBlockConfig]: """A call's layer state, per layer: the last ``att_context`` keys and values - of its attention and the last ``kernel - 1`` inputs of its depthwise conv.""" + of its attention and the last ``kernel - 1`` inputs of its depthwise conv. + 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.""" ctx = config.att_context_size[0] return { "k": RecurrentBlockConfig(shape=(config.num_heads, ctx, config.head_dim), dtype=dtype), diff --git a/mstar/model/nemotron_duplex/nemotron_duplex_model.py b/mstar/model/nemotron_duplex/nemotron_duplex_model.py index f6d8150b3e..dcafb93512 100644 --- a/mstar/model/nemotron_duplex/nemotron_duplex_model.py +++ b/mstar/model/nemotron_duplex/nemotron_duplex_model.py @@ -468,8 +468,7 @@ def _session_kv_pages(self, prefix: int, streams: int, page_size: int, sessions: return sessions * streams * -(-positions // page_size) + 1 def validate_config_yaml(self, config: dict, config_path: str) -> None: - """Refuse an autocast dtype the encoder's state is not stored in, and a - session cap that a per-session resource cannot hold. + """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``) and its nano and talker pages from @@ -477,12 +476,6 @@ def validate_config_yaml(self, config: dict, config_path: str) -> None: another ends, and its frames with it. So ``max_concurrent_requests`` sessions must fit each; the conductor queues the rest. """ - dtype, state_dtype = config.get("autocast_dtype"), str(self.get_autocast_dtype()).removeprefix("torch.") - if dtype is not None and str(dtype).removeprefix("torch.") != state_dtype: - raise ValueError( - f"{config_path}: autocast_dtype {dtype} is not the {state_dtype} that " - f"resources.{ENCODER_STATE} stores the encoder's layer state in" - ) cap = config.get("max_concurrent_requests") if cap is None: logger.warning( @@ -538,8 +531,7 @@ def get_node_resources(self) -> list[NodeResourceSpec]: resource_key=ENCODER_STATE, nodes={"conformer_encoder"}, config=RecurrentStateConfig( num_layers=self.config.stt.num_layers, - # in the dtype the layers compute it in, so a step stores what it computed - blocks=encoder_state_blocks(self.config.stt, self.get_autocast_dtype()), + blocks=encoder_state_blocks(self.config.stt), # in the engine's compute dtype max_slots=self.DEFAULT_MAMBA_SLOTS + 1, ), ), diff --git a/test/modular/test_nemotron_duplex_model.py b/test/modular/test_nemotron_duplex_model.py index 17b41b56cc..e966a7bbc4 100644 --- a/test/modular/test_nemotron_duplex_model.py +++ b/test/modular/test_nemotron_duplex_model.py @@ -171,19 +171,21 @@ def test_duplex_kv_pools_hold_the_session_cap_at_the_calls_it_serves(kwargs): "c.yaml") -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) -def test_duplex_encoder_state_is_stored_in_the_compute_dtype(dtype, monkeypatch): +@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 for the next; in another dtype than the layers compute it in, live frames drift from one-shot ones. - A deployment that computes in another dtype is refused at startup.""" + So the pool takes the dtype the engine resolved (``autocast_dtype: null`` + computes in fp32), and no deployment dtype is refused for it.""" + from mstar.engine.resources.base import EngineResourceInfo + from mstar.engine.resources.recurrent.pool import RecurrentStatePool + model = _make_model() - monkeypatch.setattr(model, "get_autocast_dtype", lambda: dtype) + model.validate_config_yaml({"autocast_dtype": autocast_dtype}, "c.yaml") (spec,) = [s for s in model.get_node_resources() if s.resource_key == ENCODER_STATE] - assert {block.dtype for block in spec.config.blocks.values()} == {dtype} - model.validate_config_yaml({"autocast_dtype": str(dtype).removeprefix("torch.")}, "c.yaml") - other = "float16" - with pytest.raises(ValueError, match=f"resources.{ENCODER_STATE}"): - model.validate_config_yaml({"autocast_dtype": other}, "c.yaml") + pool = RecurrentStatePool.build(spec, EngineResourceInfo(device=torch.device("cpu"), kv_dtype=autocast_dtype)) + want = torch.float32 if autocast_dtype is None else autocast_dtype + assert {pool.block(name, 0).dtype for name in spec.config.blocks} == {want} class _LoadReachedError(Exception): From 0164d83d0ce11b7c3b524ec474ee7d48b2112837 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 07:06:08 -0700 Subject: [PATCH 21/69] R2.5: per-message deflate is a serving setting, off in the duplex configs The deflate change turned per-message deflate off server-wide, for every WebSocket route and model, with no way back, on the premise that PCM barely compresses. As base64 JSON it does: an 80 ms speech frame deflates to 49-73% of its size (binary PCM 55-90%), digital silence to 1-2%. ws_per_message_deflate is now a serving-config key, default true (uvicorn's and main's behaviour, so no other model changes); the two Nemotron VoiceChat configs, where its event-loop cost was measured (~8% at 32 live calls on /generate/ws), set it false. docs/serving.rst documents the key with both measurements, and clients.rst says the deployment decides. Co-Authored-By: Claude Opus 5.5 --- configs/nemotron_duplex.yaml | 4 ++++ configs/nemotron_duplex_disagg.yaml | 4 ++++ docs/clients.rst | 5 +++-- docs/serving.rst | 8 ++++++++ mstar/api_server/entrypoint.py | 21 +++++++++++---------- test/modular/test_client_stream_api.py | 10 ++++++++++ test/modular/test_nemotron_duplex_model.py | 9 +++++++++ 7 files changed, 49 insertions(+), 12 deletions(-) diff --git a/configs/nemotron_duplex.yaml b/configs/nemotron_duplex.yaml index 297ad7b750..1fc4c04d3f 100644 --- a/configs/nemotron_duplex.yaml +++ b/configs/nemotron_duplex.yaml @@ -16,6 +16,10 @@ max_seq_len: 12288 # config: voicechat streaming_max_len / vllm_max_model_len # 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 99d8a58042..79f84159a6 100644 --- a/configs/nemotron_duplex_disagg.yaml +++ b/configs/nemotron_duplex_disagg.yaml @@ -21,6 +21,10 @@ max_seq_len: 12288 # 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/clients.rst b/docs/clients.rst index e553aeae89..abe906aaec 100644 --- a/docs/clients.rst +++ b/docs/clients.rst @@ -86,8 +86,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 no per-message deflate on any WebSocket route: the audio they mostly -carry barely compresses, and compressing it costs the event loop. +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 diff --git a/docs/serving.rst b/docs/serving.rst index dcbb2fd0ca..1e928de9be 100644 --- a/docs/serving.rst +++ b/docs/serving.rst @@ -274,6 +274,14 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key ``max_call_input_mb`` these count 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 39d679e045..e59584471e 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -2272,6 +2272,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. @@ -2478,16 +2488,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-wide (uvicorn negotiates it for every WebSocket - # route): /generate/ws and /v1/realtime stream audio, which - # deflate barely shrinks, and compressing it cost the event - # loop ~8% at 32 live calls on /generate/ws - ws_per_message_deflate=False, - ) - ) + server = uvicorn.Server(_uvicorn_config(app, args.host, args.port, config)) def _stop_server(): server.should_exit = True diff --git a/test/modular/test_client_stream_api.py b/test/modular/test_client_stream_api.py index 999915a9c0..5058008473 100644 --- a/test/modular/test_client_stream_api.py +++ b/test/modular/test_client_stream_api.py @@ -1598,3 +1598,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_nemotron_duplex_model.py b/test/modular/test_nemotron_duplex_model.py index e966a7bbc4..e1067b0a56 100644 --- a/test/modular/test_nemotron_duplex_model.py +++ b/test/modular/test_nemotron_duplex_model.py @@ -141,6 +141,15 @@ def test_duplex_shipped_configs_hold_their_session_cap(): 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, 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 From 98ecd5d8e9f2f3fb1df9eef84a566c81885d2f0a Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 07:07:04 -0700 Subject: [PATCH 22/69] R2.6: the streamer takes any featurizer geometry the offline encoder does Two places silently assumed the shipped featurizer (n_fft 512, hop 160): - steady_step's slices hard-code that a half window spans more than one hop and at most two; with n_fft 1024 the whole-step region was 4.3e-4 off the offline frames while the eager path matched. The streamer now says whether its geometry has a steady step; without one, is_steady is false and the encoder declares no STEP_REGION, so every step takes front + layers. - front() skipped a call whose step added no mel frame before counting its encoder frames. With a half window under a hop (n_fft 256), the end of a 1580- or 1728-sample call adds no mel frame but makes the last encoder frame final, and the call lost it (2 frames, offline 3), a regression from #594. Every busy call's frames are now counted, from its held mel frames when the step added none. Co-Authored-By: Claude Opus 5.5 --- .../nemotron_duplex/components/conformer.py | 31 +++++++++----- mstar/model/nemotron_duplex/submodules.py | 5 ++- .../test_nemotron_duplex_submodules.py | 41 +++++++++++++++++-- 3 files changed, 62 insertions(+), 15 deletions(-) diff --git a/mstar/model/nemotron_duplex/components/conformer.py b/mstar/model/nemotron_duplex/components/conformer.py index 4d81706a34..3e29ba5d99 100644 --- a/mstar/model/nemotron_duplex/components/conformer.py +++ b/mstar/model/nemotron_duplex/components/conformer.py @@ -426,6 +426,10 @@ def __init__(self, perception: "Perception"): 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) @@ -469,7 +473,7 @@ def is_steady(self, st: StreamState, chunk: torch.Tensor | None, final: 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 (not final and chunk is not None and chunk.numel() == fs + 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, frames: torch.Tensor, front: "RecurrentStatePool", @@ -540,19 +544,26 @@ def front(self, sts: list[StreamState], chunks: list[torch.Tensor | None], final 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])) - for row, (i, lo, hi) in enumerate(lo_hi): - st = sts[i] - held = min(st.mels, keep_mel) - mels = torch.cat([mel_tail[i, keep_mel - held:], mel[row, : hi - lo]]) # mels[0] is frame lo - held - st.mels = hi + 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)) - n[i] = max(0, final_encoder_frames(c, st.mels, st.samples, final[i]) - st.frames) - if n[i]: - start = max(0, f * st.frames - 2 * f) # the subsampling window's first mel frame - mel_rows.append((i, mels[start - (lo - held):], (st.frames - start // f, n[i]))) + # also without a new mel frame: the end of the call can make its last frame final + n[i] = max(0, final_encoder_frames(c, st.mels, st.samples, final[i]) - st.frames) + 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", front_rows, torch.stack(new_sig)) front.scatter_("last", front_rows, torch.stack(new_last)) front.scatter_("mel", front_rows, torch.stack(new_mel)) diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index ac90647fa3..77f418b662 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -583,10 +583,13 @@ def region(fn, inputs, keys, lease) -> PiecewiseBatchedConfig: seq_len=1, capture_batch_sizes=self.CAPTURE_BATCH_SIZES, compile=False, lease_before_step=lease, ) - return { + 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 _step_region(self, call: PiecewiseCallInputs) -> dict[str, torch.Tensor]: front, pool = call.resources[ENCODER_FRONT], call.resources[ENCODER_STATE] diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index d88ce84aa3..86b0553923 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -541,14 +541,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(): @@ -562,8 +562,8 @@ class _Encoder: """The encoder node with its state pool, driven as the engine drives an eager step: declare, admit and plan the pool, preprocess, forward, commit.""" - def __init__(self, max_session_samples=16000 * 10): - p = _small_perception() + def __init__(self, max_session_samples=16000 * 10, **featurizer): + p = _small_perception(**featurizer) self.node = ConformerEncoderSubmodule(p, nn.Identity(), nn.Identity(), config=NemotronDuplexConfig(), max_session_samples=max_session_samples) def pool(key, num_layers, blocks): @@ -837,6 +837,39 @@ def test_encoder_and_codec_take_no_batch_wider_than_their_graphs(): 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_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.""" From fc5db609ded8d7a3c20fe55a7021150d3f604890 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 07:08:18 -0700 Subject: [PATCH 23/69] R2.7: refuse a checkpoint the declared resources were not sized for The resources are declared from the dataclass defaults (get_node_resources runs before the checkpoint is read), while the encoder, talker and codec modules are built from the checkpoint's config.json; the codec's region was sized from the defaults too. A checkpoint with another encoder depth or codebook count would have run against pools and buffers of another shape. When the checkpoint's config is read, the model now declares its resources from it as well and refuses a checkpoint that sizes any one otherwise, naming it. The codec takes the checkpoint's config, as the talker does, since it decodes the codes the talker emits. The released checkpoint's config.json passes the check (its talker sampling knobs differ from the defaults; no resource does). Co-Authored-By: Claude Opus 5.5 --- .../nemotron_duplex/nemotron_duplex_model.py | 23 ++++++++++++++--- test/modular/test_nemotron_duplex_model.py | 25 +++++++++++++++++++ 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/mstar/model/nemotron_duplex/nemotron_duplex_model.py b/mstar/model/nemotron_duplex/nemotron_duplex_model.py index dcafb93512..44eda9fab4 100644 --- a/mstar/model/nemotron_duplex/nemotron_duplex_model.py +++ b/mstar/model/nemotron_duplex/nemotron_duplex_model.py @@ -44,6 +44,7 @@ """ from __future__ import annotations +import copy import logging from pathlib import Path @@ -1223,12 +1224,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. @@ -1307,4 +1323,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/test/modular/test_nemotron_duplex_model.py b/test/modular/test_nemotron_duplex_model.py index e1067b0a56..62ca919792 100644 --- a/test/modular/test_nemotron_duplex_model.py +++ b/test/modular/test_nemotron_duplex_model.py @@ -227,6 +227,31 @@ def built(self, module, device, autocast_dtype, local_dir): 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_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) From a322b955e3fb7db6a2b2c2c38cb71838974f0b46 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 07:08:29 -0700 Subject: [PATCH 24/69] R2.8: tell live clients the whole-step graph keeps to the call's frame grid A live step replays the encoder's whole-step graph only when the call's audio so far is a whole number of 80 ms frames, counted from its start, so a call that opens with a one-second clip (16000 % 1280 = 640) never takes it, while clients.rst read as if any run of 80 ms appends would. clients.rst and models.rst now tell clients to keep the opening clip and every append a multiple of 1280 samples. A phase-general steady step would need its mel and subsampling windows to shift with the phase; not done. Co-Authored-By: Claude Opus 5.5 --- docs/clients.rst | 7 +++++-- docs/models.rst | 11 ++++++----- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/docs/clients.rst b/docs/clients.rst index abe906aaec..44fffd7240 100644 --- a/docs/clients.rst +++ b/docs/clients.rst @@ -181,8 +181,11 @@ plus one each for the opening message's media and the end of input; the next end with ``status`` 503. 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 a call whose appends are exactly 80 ms (1280 samples at 16 kHz) on its -frame grid; other sizes are served the same, more slowly (see :doc:`models`). +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 1d06439345..18e370c812 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -534,11 +534,12 @@ Nemotron VoiceChat (``nemotron_duplex``) notes session each: its layer state (attention keys / values and conv inputs, ``encoder_state``) and its newest audio and mel frames (``encoder_front``). 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; other calls - take the front end eagerly, then the layer stack as a graph when each got at - most one frame, or eagerly (a one-shot upload, a large append, the call's - end). Appends of other sizes are served the same, without the whole-step - graph. + 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, or + eagerly (a one-shot upload, a large append, the call's end). 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 From ba262f7cfd336ae8e4c32e62c2c55e9cc0d48db2 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 07:09:09 -0700 Subject: [PATCH 25/69] R2.9: decide once whether a call takes the whole-step region declare_step (to leave both pools to STEP_REGION) and forward_batched (to route the calls to it) each decided a call's steadiness from its stream state, tied only by an assert that python -O drops. prepare_inputs now decides it, once per call and step, and passes it in the call's kwargs; the declaration and the forward read that. Declaring only the eager rows in a mixed step (so the step region's calls are not planned twice) is left out: when the region cannot run a sub-batch the forward falls back to the eager path, which then needs those rows planned. Co-Authored-By: Claude Opus 5.5 --- mstar/model/nemotron_duplex/submodules.py | 24 ++++++++----------- .../test_nemotron_duplex_submodules.py | 18 ++++++++++++++ 2 files changed, 28 insertions(+), 14 deletions(-) diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index 77f418b662..ade843aff7 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -497,15 +497,19 @@ 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) return NodeInputs( tensor_inputs={} if wav is None else {"chunk": wav}, - kwargs={"final": bool(is_final_stream_chunk)}, + # whether the whole-step region serves the call: decided here, once, + # for the declaration and the forward to share + kwargs={"final": final, "steady": st is not None and self.streamer.is_steady(st, wav, final)}, ) 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], } def can_batch(self, batch, model_inputs) -> bool: @@ -530,9 +534,7 @@ def declare_step(self, graph_walk, request_ids, inputs, slot_lease=None, piecewi # over the whole batch, at the bucket its lease names, and the region's # own step owns both pools. 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( - self._is_steady(rid, inp) for rid, inp in zip(request_ids, inputs, strict=True) - ): + 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 the # eager path reads that plan. A region the forward still replays for @@ -542,11 +544,6 @@ def declare_step(self, graph_walk, request_ids, inputs, slot_lease=None, piecewi # can't refuse a call that has no slot yet. return self._declare(request_ids, (ENCODER_FRONT, ENCODER_STATE)) - def _is_steady(self, rid, inputs: NodeInputs) -> bool: - st = self.request_state(rid).get(self.STATE_KEY) - return st is not None and self.streamer.is_steady(st, inputs.tensor_inputs.get("chunk"), - inputs.kwargs["final"]) - @property def streamer(self) -> ConformerStreamer: if self._streamer is None: @@ -623,7 +620,7 @@ def _counts(values, device) -> torch.Tensor: t = torch.tensor(values, dtype=torch.long) return (t.pin_memory() if device.type == "cuda" else t).to(device, non_blocking=True) - def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, **kwargs): + def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, steady=None, **kwargs): rids = list(engine_inputs.request_ids) device = self.perception.proj.weight.device streamer = self.streamer @@ -635,8 +632,7 @@ def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, **k # calls one frame of audio further on the frame grid: the whole step as one graph, run last step_runner = runners.get(self.STEP_REGION) - steady = [i for i, (st, c, f) in enumerate(zip(states, chunks, final, strict=True)) - if streamer.is_steady(st, c, f)] if step_runner is not None else [] + steady = [i for i, s in enumerate(steady) if s] if step_runner is not None else [] if engine_inputs.step is None: # declare_step left both pools to the region (see there) assert len(steady) == len(rids), "an encoder step declared nothing but is not all steady" @@ -708,10 +704,10 @@ def _eager_layers(self, enc, pool, rows, device) -> list[torch.Tensor]: y = self.streamer.layers(enc.x, pool, rows, counts[0], counts[1]) return [row[:n] for row, n in zip(y, enc.n, strict=True)] - def forward(self, graph_walk, engine_inputs, chunk=None, final=None, **kwargs) -> NameToTensorList: + 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): diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index 86b0553923..45d125052d 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -870,6 +870,24 @@ def test_encoder_ends_a_call_with_its_last_frame_when_the_end_adds_no_mel_frame( 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_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.""" From 1d09a92b2dc8948be50e0c0b60a29b77e3d256e9 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 07:09:50 -0700 Subject: [PATCH 26/69] R2.10: an encoder step copies its calls' chunks to the device at once Each call's chunk arrives on the host and was copied to the device on its own, from pageable memory, so a 32-call step made 32 blocking transfers; the eager subsampling's per-row lengths were one more. The step now concatenates its host chunks, pins them once and copies them in one transfer, then splits them on the device; the lengths go through the same pinned copy as the step's other counts (device_counts, moved next to the streamer that needs it). Unmeasured on a GPU. Co-Authored-By: Claude Opus 5.5 --- .../nemotron_duplex/components/conformer.py | 9 +++++- mstar/model/nemotron_duplex/submodules.py | 28 ++++++++++++------- .../test_nemotron_duplex_submodules.py | 14 ++++++++++ 3 files changed, 40 insertions(+), 11 deletions(-) diff --git a/mstar/model/nemotron_duplex/components/conformer.py b/mstar/model/nemotron_duplex/components/conformer.py index 3e29ba5d99..8b5c20db54 100644 --- a/mstar/model/nemotron_duplex/components/conformer.py +++ b/mstar/model/nemotron_duplex/components/conformer.py @@ -390,6 +390,13 @@ 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 frames before the step and new ones.""" @@ -459,7 +466,7 @@ def _subsample(self, windows: torch.Tensor, lens: list[int] | None = None) -> to 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 = torch.tensor(lens, pin_memory=out.is_cuda).to(out.device, non_blocking=True).view(-1, 1, 1, 1) + 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) diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index ade843aff7..f2bbb61079 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -47,7 +47,7 @@ ) 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, @@ -614,18 +614,26 @@ def _pool(engine_inputs, key) -> tuple: return pool, pool.addressing("main") @staticmethod - def _counts(values, 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) + 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, **kwargs): rids = list(engine_inputs.request_ids) device = self.perception.proj.weight.device streamer = self.streamer states = [self._stream_state(rid) for rid in rids] - chunks = [c.to(device, non_blocking=True) if c is not None else None for c in chunk] + chunks = self._to_device(chunk, device) final = list(final) runners = engine_inputs.piecewise_runners or {} frames: list[torch.Tensor | None] = [None] * len(rids) @@ -666,7 +674,7 @@ def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, ste new = [i for i, n in enumerate(enc.n) if n] if runner is not None and enc.x.shape[1] == 1 and runner.can_run(len(new)): # only the calls with a frame: the region plans its own rows - counts = self._counts([[enc.frames[i] for i in new], [1] * len(new)], device) + counts = device_counts([[enc.frames[i] for i in new], [1] * len(new)], device) y = runner.run(static_inputs={"x": torch.stack([enc.x[i] for i in new]), "frames": counts[0], "n": counts[1]}, request_ids=[rids[i] for i in new])["y"] @@ -680,7 +688,7 @@ def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, ste if steady: y = step_runner.run( static_inputs={"chunk": torch.stack(steady_chunks), - "frames": self._counts([states[i].frames for i in steady], device)}, + "frames": device_counts([states[i].frames for i in steady], device)}, request_ids=[rids[i] for i in steady], )["y"] for row, i in enumerate(steady): @@ -700,7 +708,7 @@ def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, ste def _eager_layers(self, enc, pool, rows, device) -> list[torch.Tensor]: """The layer stack eagerly over ``enc`` at ``rows``: each call's new frames.""" - counts = self._counts([enc.frames, enc.n], device) + counts = device_counts([enc.frames, enc.n], device) y = self.streamer.layers(enc.x, pool, rows, counts[0], counts[1]) return [row[:n] for row, n in zip(y, enc.n, strict=True)] diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index 45d125052d..d5a19d86f0 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -888,6 +888,20 @@ def decided_already(*args): 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_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.""" From 9306c94f41de11f3cf8246a084820bddd7bf54fb Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 07:10:21 -0700 Subject: [PATCH 27/69] R2.13: the encoder tests take their region leases from the engine's own method The test harness re-implemented Engine._maybe_lease_piecewise_regions to decide which regions lease before the declaration, so a change to the engine's rule would not have reached the tests. The harness now calls the real method, on an engine holding only the node's piecewise runners; the comment that still called those runners stand-ins goes. Co-Authored-By: Claude Opus 5.5 --- test/modular/test_nemotron_duplex_submodules.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index d5a19d86f0..8c8f762434 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -575,7 +575,7 @@ def pool(key, num_layers, blocks): self.front = pool(ENCODER_FRONT, 1, front_state_blocks(p.config)) self.resources = {ENCODER_STATE: self.pool, ENCODER_FRONT: self.front} self.runner = StepRunner(self.resources) - self.regions = {} # stand-ins for the piecewise runners, by label, when a test sets them + 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": []} @@ -584,9 +584,13 @@ def prepare(self, handle, chunk, final=False): def declare(self, rids, inputs): """The node's declaration, with the leases the engine takes for it first.""" - leases = {label: lease for label, region in self.regions.items() - if region.lease_before_step and (lease := region.lease_slot(len(inputs))) is not None} - return self.node.declare_step("encode", rids, inputs, piecewise_leases=leases) + 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: for rid in rids: From bca60c8de00c65d635c3658cdc165471ce492757 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 07:10:39 -0700 Subject: [PATCH 28/69] R2.14: replay the encoder's step region and the codec's region with a padding row No test replayed STEP_REGION or the codec's decode region in a bucket wider than its batch: the encoder's live-step test put two calls in a bucket of two, the codec's ran one call. Three calls one frame on now replay the step region in a bucket of four, and three calls in step decode their steady windows through the codec region in a bucket of four; each call still gets its offline frames / full-history decode. Co-Authored-By: Claude Opus 5.5 --- .../test_nemotron_duplex_submodules.py | 56 +++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index 8c8f762434..cffe058043 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -476,6 +476,46 @@ def test_codec_stream_is_the_full_history_decode_with_the_real_decoder(frames, r 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 @@ -906,6 +946,22 @@ def test_encoder_moves_a_steps_chunks_to_the_device_in_one_copy(): ("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.""" From 6737e6db8e5fba8c1dccd171231d39ae3be4b0c9 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 07:11:26 -0700 Subject: [PATCH 29/69] R2.16: the default session count is its own constant, not the capture list's DEFAULT_MAMBA_SLOTS, the session count every per-session pool is sized for, was SESSION_BATCH_SIZES[-1], so widening the capture list would have resized every pool. It is now DEFAULT_SESSIONS = 32, named for what it counts (it sizes the encoder pools and both KV pools too, not only mamba_state). The shipped configs' comments name it and nano_kv with the other pools. Co-Authored-By: Claude Opus 5.5 --- configs/nemotron_duplex.yaml | 5 ++-- configs/nemotron_duplex_disagg.yaml | 5 ++-- mstar/model/nemotron_duplex/config.py | 2 +- .../nemotron_duplex/nemotron_duplex_model.py | 27 +++++++++---------- test/modular/test_nemotron_duplex_model.py | 4 +-- .../test_nemotron_duplex_submodules.py | 2 +- 6 files changed, 23 insertions(+), 22 deletions(-) diff --git a/configs/nemotron_duplex.yaml b/configs/nemotron_duplex.yaml index 1fc4c04d3f..80da7e1e8d 100644 --- a/configs/nemotron_duplex.yaml +++ b/configs/nemotron_duplex.yaml @@ -11,8 +11,9 @@ max_seq_len: 12288 # config: voicechat streaming_max_len / vllm_max_model_len # Sessions at once; the overflow queues. Each session holds a slot in # resources mamba_state, encoder_state and encoder_front and its pages in -# talker_kv, all sized for 32 (DEFAULT_MAMBA_SLOTS); to raise the cap, raise -# their max_slots (cap + 1, for the sink) and talker_kv's max_num_pages with it. +# nano_kv and talker_kv, all sized for 32 (DEFAULT_SESSIONS); to raise the cap, +# raise their max_slots (cap + 1, for the sink) 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 diff --git a/configs/nemotron_duplex_disagg.yaml b/configs/nemotron_duplex_disagg.yaml index 79f84159a6..7eb234375c 100644 --- a/configs/nemotron_duplex_disagg.yaml +++ b/configs/nemotron_duplex_disagg.yaml @@ -16,8 +16,9 @@ max_seq_len: 12288 # Sessions at once; the overflow queues. Each session holds a slot in # resources mamba_state, encoder_state and encoder_front and its pages in -# talker_kv, all sized for 32 (DEFAULT_MAMBA_SLOTS); to raise the cap, raise -# their max_slots (cap + 1, for the sink) and talker_kv's max_num_pages with it. +# nano_kv and talker_kv, all sized for 32 (DEFAULT_SESSIONS); to raise the cap, +# raise their max_slots (cap + 1, for the sink) 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 diff --git a/mstar/model/nemotron_duplex/config.py b/mstar/model/nemotron_duplex/config.py index 5c88760096..6f028d6fb6 100644 --- a/mstar/model/nemotron_duplex/config.py +++ b/mstar/model/nemotron_duplex/config.py @@ -34,7 +34,7 @@ ENCODER_FRONT = "encoder_front" # recurrent-state pool: each call's streaming-encoder audio and mel tails # Every node's CUDA-graph batch sizes: a row per live session, up to the sessions -# a deployment serves by default (``NemotronDuplexModel.DEFAULT_MAMBA_SLOTS``). +# a deployment serves by default (``NemotronDuplexModel.DEFAULT_SESSIONS``). SESSION_BATCH_SIZES = [1, 2, 4, 8, 16, 32] diff --git a/mstar/model/nemotron_duplex/nemotron_duplex_model.py b/mstar/model/nemotron_duplex/nemotron_duplex_model.py index 44eda9fab4..5027d7ecb7 100644 --- a/mstar/model/nemotron_duplex/nemotron_duplex_model.py +++ b/mstar/model/nemotron_duplex/nemotron_duplex_model.py @@ -83,7 +83,6 @@ NANO_ATTN, NANO_KV, NANO_SAMPLER, - SESSION_BATCH_SIZES, TALKER_ATTN, TALKER_KV, TALKER_POS, @@ -417,13 +416,13 @@ def _talker_tokenizer(self): # ------------------------------------------------------------------- # Sessions resident at once (a slot each in mamba_state, ~137 MB in fp32 for - # Nemotron-H, and in encoder_state / encoder_front; pages in talker_kv). A - # deployment raising it raises all four (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_MAMBA_SLOTS = SESSION_BATCH_SIZES[-1] # 32 sessions; padding rows take no slot (#258) + # Nemotron-H, and in encoder_state / encoder_front; pages in nano_kv and + # talker_kv). A deployment raising it raises all five (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 @@ -455,15 +454,15 @@ def max_session_frames(self) -> int: return self.config.stt.num_frames(self.max_session_samples) def talker_kv_pages(self, page_size: int, sessions: int | None = None) -> int: - """Pages for ``sessions`` (DEFAULT_MAMBA_SLOTS) sessions of max_session_frames, both CFG streams.""" + """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_MAMBA_SLOTS) sessions of the longest prompt and max_session_frames.""" + """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_MAMBA_SLOTS if sessions is None else sessions + 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 sessions * streams * -(-positions // page_size) + 1 @@ -533,14 +532,14 @@ def get_node_resources(self) -> list[NodeResourceSpec]: 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_MAMBA_SLOTS + 1, + 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_MAMBA_SLOTS + 1, + max_slots=self.DEFAULT_SESSIONS + 1, ), ), KVSpec(resource_key=NANO_KV, nodes={"nano_llm"}, config=kv_config), @@ -559,8 +558,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( diff --git a/test/modular/test_nemotron_duplex_model.py b/test/modular/test_nemotron_duplex_model.py index 62ca919792..caff58ff7b 100644 --- a/test/modular/test_nemotron_duplex_model.py +++ b/test/modular/test_nemotron_duplex_model.py @@ -91,7 +91,7 @@ def test_duplex_node_resources(): assert isinstance(enc, RecurrentStateSpec) and enc.config.num_layers == stt.num_layers assert enc.config.blocks["k"].shape == (stt.num_heads, stt.att_context_size[0], stt.head_dim) assert enc.config.blocks["conv"].shape == (stt.d_model, stt.conv_kernel_size - 1) - assert enc.config.usable_slots == model.DEFAULT_MAMBA_SLOTS + assert enc.config.usable_slots == model.DEFAULT_SESSIONS eartts = model.config.eartts tkv = by_key[TALKER_KV] @@ -113,7 +113,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} diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index cffe058043..876e21b28a 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -205,7 +205,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) From 6304de47e2824a1deae9e8b74c576bd873d5cff4 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 07:12:07 -0700 Subject: [PATCH 30/69] R2.17: correct the duplex docs and comments the review found stale - models.rst offered mamba_state's state_dtype, which its pool refuses (a Mamba-2 pool has no "state" block); max_slots is its only knob. - models.rst and the encoder's docstring listed the call's end among the eager layer passes; it usually releases one frame and replays the layer graph. An upload or large append runs in a pass of its own, not a step. - The codec's "fp32 codec" comment: with autocast off the codec runs in the dtype it was loaded in, the engine's (bf16 by default). - The nano's capture comment said padding rows take a mamba slot transiently; they take none. Co-Authored-By: Claude Opus 5.5 --- docs/models.rst | 9 +++++---- mstar/model/nemotron_duplex/submodules.py | 11 ++++++----- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/docs/models.rst b/docs/models.rst index 18e370c812..9d1321271a 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -438,7 +438,7 @@ 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 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 @@ -537,9 +537,10 @@ Nemotron VoiceChat (``nemotron_duplex``) notes 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, or - eagerly (a one-shot upload, a large append, the call's end). Appends of other - sizes are served the same, without the whole-step graph. + 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, + run in a pass of 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/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index f2bbb61079..a113782568 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -106,8 +106,8 @@ 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. + # 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 @@ -446,8 +446,8 @@ class ConformerEncoderSubmodule(NodeSubmodule): 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, else eagerly (a - one-shot upload, a large append, the call's end). + (``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). Holds the perception stack plus the RNN-T decoder/joint (loaded; the user-transcript channel is not emitted yet). @@ -971,7 +971,8 @@ class AudioCodecDecoderSubmodule(NodeSubmodule): live load runs many small decodes. Declares no resources. """ - # Per-request context is keyed by request id -> eager around the decode; 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 From 26d62214921dfa4b905f62183175ce6a76eec2b6 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Mon, 19 Oct 2026 14:15:57 -0700 Subject: [PATCH 31/69] R1.9: keep cuFFT out of the captured encoder step and codec decode conformer_step captured torch.stft (steady_step's mel) and codec_decode captured torch.fft.irfft (the codec's iSTFT). torch's cuFFT plan cache is a per-device LRU whose eviction calls cufftDestroy, and graph replays don't touch the LRU, so enough newer eager FFT geometries (eager front-end windows, uploads, eager codec windows) destroy the plans the graphs replay. Both regions now take their fixed-size transform as a float64 matmul against a precomputed real-DFT basis (components/dft.py): the STFT basis folds in the checkpoint's window and is built when the streamer is (after load, before capture); the n=16 inverse basis is cached per device and first built in the region's eager warmup. float64 keeps TF32 and autocast out of it. The eager paths keep torch.stft / irfft, so one-shot output stays bitwise #594's and main's Perception, and the eager streamer stays bitwise #594's. GPU, 1x H100 PCIe, served checkpoint, regions captured as the engine does: - after cufft_plan_cache.clear() + eager FFT churn + poisoned free memory, the #595 head hits an illegal memory access on its next conformer_step replay; with this change encoder frames (2 clips, single and batched) and codec windows (bs 32) are bitwise unchanged, fp32 and bf16. - matmul vs fp64 reference: STFT within fp32 rounding (torch.stft fp32: 1e-6 on values up to 16); irfft 3e-8 (torch fp32: 2.4e-7); same under TF32 and bf16 autocast. - encoder, captured vs main offline: fp32 min cos 0.9999996 (C1.2 bar 0.9999); bf16 mean|diff| 0.99-1.01x the bf16-vs-fp32 yardstick (bar 2x). Captured vs eager in fp32: max 4.7e-7 (was 2.3e-7 through cuFFT). Codec captured vs eager: 1.8e-7, at most one int16 step. - graph replay at bucket 32: conformer_step 8209 -> 8226 us (bf16), codec_decode 18657 -> 18669 us (bf16): +0.2% / +0.1%. Out of scope: an engine-level guard against FFTs in captures; chatterbox's captured vocoder on main has the same exposure. Co-Authored-By: Claude Opus 5.5 --- .../nemotron_duplex/components/audio_codec.py | 24 ++-- .../nemotron_duplex/components/conformer.py | 22 +++- mstar/model/nemotron_duplex/components/dft.py | 75 ++++++++++++ mstar/model/nemotron_duplex/submodules.py | 3 +- .../test_nemotron_duplex_submodules.py | 113 ++++++++++++++++++ 5 files changed, 223 insertions(+), 14 deletions(-) create mode 100644 mstar/model/nemotron_duplex/components/dft.py diff --git a/mstar/model/nemotron_duplex/components/audio_codec.py b/mstar/model/nemotron_duplex/components/audio_codec.py index 9c006e3d0b..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: @@ -238,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): @@ -256,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 @@ -274,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 8b5c20db54..acb668633c 100644 --- a/mstar/model/nemotron_duplex/components/conformer.py +++ b/mstar/model/nemotron_duplex/components/conformer.py @@ -27,6 +27,7 @@ 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: @@ -442,15 +443,22 @@ def __init__(self, perception: "Perception"): 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 _mel(self, batch: torch.Tensor) -> torch.Tensor: - """``(B, samples)`` preemphasized windows -> ``(B, T, n_mels)`` log-mel frames.""" + 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 - spec = 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) - spec = torch.view_as_real(spec) + 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) @@ -488,6 +496,8 @@ def steady_step(self, chunk: torch.Tensor, frames: torch.Tensor, front: "Recurre rows: "RecurrentAddressing") -> 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 @@ -501,7 +511,7 @@ def steady_step(self, chunk: torch.Tensor, frames: torch.Tensor, front: "Recurre 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]) + 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", front_rows, y) 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/submodules.py b/mstar/model/nemotron_duplex/submodules.py index a113782568..841e98fba3 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -1048,7 +1048,8 @@ def make_static_inputs(shape: PiecewiseCaptureShape) -> dict[str, torch.Tensor]: } def _decode_region(self, call: PiecewiseCallInputs) -> dict[str, torch.Tensor]: - audio, _ = self.codec.decode(call.static_inputs["codes"]) + # 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]: diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index 876e21b28a..b04afa1815 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -1026,3 +1026,116 @@ def test_encoder_state_is_per_call_and_cleaned_up(): pool.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 + + +# -- 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() From 1873ecb9757f2b8d646da4f086c72b4b793046b1 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 00:44:54 -0700 Subject: [PATCH 32/69] nemotron_duplex: the encoder no longer takes max_session_samples With audio_in's max_items enforcing the session cap in the client-stream layer, the encoder's own check is gone, and since its state moved into bounded pools nothing else in it reads the session length. Drop the parameter, its plumbing in the model and the tests' settings of it. Co-Authored-By: Claude Opus 5.5 --- mstar/model/nemotron_duplex/nemotron_duplex_model.py | 5 +---- mstar/model/nemotron_duplex/submodules.py | 3 +-- test/modular/test_nemotron_duplex_submodules.py | 8 +++----- 3 files changed, 5 insertions(+), 11 deletions(-) diff --git a/mstar/model/nemotron_duplex/nemotron_duplex_model.py b/mstar/model/nemotron_duplex/nemotron_duplex_model.py index 5027d7ecb7..ca574f5e26 100644 --- a/mstar/model/nemotron_duplex/nemotron_duplex_model.py +++ b/mstar/model/nemotron_duplex/nemotron_duplex_model.py @@ -1286,10 +1286,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, diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index 841e98fba3..48ff6ab0a5 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -467,13 +467,12 @@ class ConformerEncoderSubmodule(NodeSubmodule): 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( diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index b04afa1815..cab16af198 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -602,10 +602,9 @@ class _Encoder: """The encoder node with its state pool, driven as the engine drives an eager step: declare, admit and plan the pool, preprocess, forward, commit.""" - def __init__(self, max_session_samples=16000 * 10, **featurizer): + def __init__(self, **featurizer): p = _small_perception(**featurizer) - self.node = ConformerEncoderSubmodule(p, nn.Identity(), nn.Identity(), config=NemotronDuplexConfig(), - max_session_samples=max_session_samples) + 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)) @@ -822,7 +821,6 @@ def test_encoder_a_calls_last_frames_do_not_depend_on_its_batch(b_len, a_chunk, 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.node.max_session_samples = a.numel() got_b = [] for k, pb in enumerate(starts): b_final = b_end_rides and k == len(starts) - 1 @@ -871,7 +869,7 @@ def test_encoder_and_codec_take_no_batch_wider_than_their_graphs(): nodes = { "conformer_encoder": ConformerEncoderSubmodule(nn.Identity(), nn.Identity(), nn.Identity(), - config=NemotronDuplexConfig(), max_session_samples=16000), + config=NemotronDuplexConfig()), "audio_codec": AudioCodecDecoderSubmodule(codec=nn.Identity(), config=NemotronDuplexConfig()), } engine = Engine.__new__(Engine) From 7497253071b56fa5ece0852277aa31ab6a1bb696 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 00:47:10 -0700 Subject: [PATCH 33/69] nemotron_duplex: separate long encoder chunks by can_batch; drop RecurrentAddressing.narrow ConformerEncoderSubmodule.forward_batched split its own batch: a chunk over BATCHED_CHUNK_S (1 s; an upload or a large append) ran alone inside the forward, on a narrow of the step's plan, because a pass pads every call to its longest chunk. Batch composition is the scheduler's and the engine's business, not the forward's (AGENTS.md invariant 1). prepare_inputs now reports the chunk's length as input_seq_len (in samples), and can_batch refuses a batch that holds a chunk over a second alongside other calls. The engine then takes its per-request path for that step: each call declares, admits, plans, runs and commits on its own (Engine._exec_per_request), so the long chunk runs as a pass of its own and the other calls still run in the same step, none starved. The micro-scheduler itself consults neither can_batch nor input_seq_len: it batches by readiness, caps and capture groups, and splitting by capture key needs a whole-forward capture, which the encoder (piecewise regions only) has none of. The cost is that a step holding such a chunk runs its other calls one at a time; it happens only on uploads and large appends. The alternative, cutting a long upload into items of a second or less in process_prompt, was rejected: it changes one-shot output (a one-file upload is one encoder pass today, bit for bit the offline encoder's; a different chunking moves bf16 frames by about 1%) and turns a 160 s upload into 160 encoder steps. The math for every call is the one the forward ran before: the long chunk's pass is the same streamer front and eager layer stack, now over its own step's plan instead of a narrow of a shared one. With nothing left that runs some rows of a plan apart, RecurrentAddressing.narrow, its test and its mention in docs/adding_models.rst are gone. Tests: the encoder harness runs a step the way the engine does, per call when can_batch refuses it; the long-chunk test checks input_seq_len, can_batch, that each call of such a step runs alone and that calls of a second or less share a pass, and the offline parity of every call. The parity tests that mix a long chunk with others now go through the per-call path. Co-Authored-By: Claude Opus 5.5 --- docs/adding_models.rst | 3 +- docs/models.rst | 5 +- mstar/engine/resources/recurrent/pool.py | 14 ----- mstar/model/nemotron_duplex/submodules.py | 33 +++++------- test/modular/test_nemotron_duplex_model.py | 2 +- .../test_nemotron_duplex_submodules.py | 51 +++++++++++++------ test/modular/test_recurrent_padding_rows.py | 21 -------- 7 files changed, 53 insertions(+), 76 deletions(-) diff --git a/docs/adding_models.rst b/docs/adding_models.rst index ffa879d2dc..3b5daf950c 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -442,8 +442,7 @@ The spec types are: caches store); ``max_slots`` (one more than the requests, for the sink that padding rows address) is deployment-tunable. A recurrent kernel reads the step's ``addressing``; plain torch reads and writes a block through - ``gather`` / ``scatter_`` with that addressing (or its ``narrow``, for rows a - forward runs apart from the rest), so eager and captured forwards + ``gather`` / ``scatter_`` with that addressing, so eager and captured forwards address the same slots. Those two need the sink: a pool built with ``disable_sink_slot`` refuses them for a padded step. * - ``SamplerSpec(vocab_size=..., enable_repetion_penalty=...)`` diff --git a/docs/models.rst b/docs/models.rst index 9d1321271a..b2c5bfbb92 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -538,8 +538,9 @@ Nemotron VoiceChat (``nemotron_duplex``) notes 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, - run in a pass of its own). Appends of other sizes are served the same, + 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 diff --git a/mstar/engine/resources/recurrent/pool.py b/mstar/engine/resources/recurrent/pool.py index 1078d88d8c..581aeba235 100644 --- a/mstar/engine/resources/recurrent/pool.py +++ b/mstar/engine/resources/recurrent/pool.py @@ -65,20 +65,6 @@ class RecurrentAddressing: host_slot_indices: Sequence[int] | None = None host_has_state: Sequence[bool] | None = None - def narrow(self, start: int, length: int = 1) -> RecurrentAddressing: - """Rows ``start .. start + length`` of the step, for a forward that runs - them apart from the rest (a long row that would pad the others). Views - of this plan: valid exactly as long as it is.""" - end = start + length - assert 0 <= start < end <= self.num_rows, f"rows {start}..{end} of a step of {self.num_rows}" - return RecurrentAddressing( - slot_indices=self.slot_indices[start:end], - has_state=self.has_state[start:end], - num_rows=length, - host_slot_indices=None if self.host_slot_indices is None else self.host_slot_indices[start:end], - host_has_state=None if self.host_has_state is None else self.host_has_state[start:end], - ) - class RecurrentStatePool(Resource): @classmethod diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index 48ff6ab0a5..c8c153e0f7 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -436,9 +436,9 @@ 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, on its own row of the step's - plan. A call's state is bounded whatever its length: its newest audio and + 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), and the last 70 frames of each layer's attention keys and values plus the conv tails (``ENCODER_STATE``); its host counters live in its per-request state. @@ -502,6 +502,7 @@ def prepare_inputs( # whether the whole-step region serves the call: decided here, once, # for the declaration and the forward to share kwargs={"final": final, "steady": st is not None and self.streamer.is_steady(st, wav, final)}, + input_seq_len=0 if wav is None else wav.numel(), # samples, for can_batch ) def preprocess(self, graph_walk, engine_inputs, inputs): @@ -512,7 +513,12 @@ def preprocess(self, graph_walk, engine_inputs, 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 @@ -633,7 +639,6 @@ def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, ste streamer = self.streamer states = [self._stream_state(rid) for rid in rids] chunks = self._to_device(chunk, device) - final = list(final) runners = engine_inputs.piecewise_runners or {} frames: list[torch.Tensor | None] = [None] * len(rids) @@ -650,21 +655,9 @@ def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, ste chunks[i] = None # 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. A pass pads every call to - # its longest chunk, so a chunk of more than a second (an upload, a large - # append) runs alone, on its own row of the plan, and the batched pass - # leaves that call as it is. - longest = int(self.BATCHED_CHUNK_S * self.perception.config.sample_rate) - for i, c in enumerate(chunks): - if c is not None and c.numel() > longest: - front, front_rows = self._pool(engine_inputs, ENCODER_FRONT) - enc = streamer.front([states[i]], [c], [final[i]], front, front_rows.narrow(i)) - if enc is not None: - pool, rows = self._pool(engine_inputs, ENCODER_STATE) - frames[i] = self._eager_layers(enc, pool, rows.narrow(i), device)[0] - streamer.advance([states[i]], enc) - chunks[i], final[i] = None, False - # the batched pass: its layers as a graph when no call got more than a frame + # 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, *self._pool(engine_inputs, ENCODER_FRONT)) diff --git a/test/modular/test_nemotron_duplex_model.py b/test/modular/test_nemotron_duplex_model.py index caff58ff7b..8715a5b741 100644 --- a/test/modular/test_nemotron_duplex_model.py +++ b/test/modular/test_nemotron_duplex_model.py @@ -503,7 +503,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 cab16af198..5a98e0f20d 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -632,6 +632,16 @@ def declare(self, rids, 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: self.pool.ingest_request(rid) self.front.ingest_request(rid) @@ -770,7 +780,7 @@ def test_encoder_paths_read_the_plan_current_when_they_use_it(monkeypatch): """Every read and write of a pool's state, eager or in a region, uses the plan the pool holds at that moment, never one held across another plan: the eager path runs on this step's plan before a region plans its own. A - call with a chunk over a second runs on a narrow of that plan.""" + 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) for pool in enc.resources.values(): @@ -833,33 +843,42 @@ def test_encoder_a_calls_last_frames_do_not_depend_on_its_batch(b_len, a_chunk, 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 pass pads every call to its longest chunk: an upload or a large append - runs alone, on its own row of the step's plan, and only the calls with a - second or less share a pass, whose layers then take at most a second of - frames per call.""" +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, layers = enc.node.streamer, [], [] - front, layer_stack = streamer.front, streamer.layers + 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) - def spy_layers(x, pool, rows, *rest): - layers.append((x.shape[0], x.shape[1], rows.num_rows)) - return layer_stack(x, pool, rows, *rest) - - streamer.front, streamer.layers = spy_front, spy_layers + 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} + 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]] - assert all(b == rows for b, _t, rows in layers) - assert max(t for b, t, _rows in layers if b > 1) == len(out[4]) # the widest call of a second or less + 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[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 diff --git a/test/modular/test_recurrent_padding_rows.py b/test/modular/test_recurrent_padding_rows.py index 391f2925ee..cd9887b923 100644 --- a/test/modular/test_recurrent_padding_rows.py +++ b/test/modular/test_recurrent_padding_rows.py @@ -211,24 +211,3 @@ def test_gather_reads_with_the_planned_indices_as_they_are(): pool.gather("conv", rows) ops = {event.name for event in prof.events()} assert "aten::index_select" in ops and not ops & {"aten::to", "aten::_to_copy"}, ops - - -def test_narrow_addresses_some_rows_of_the_step(): - """``narrow`` is a run of the step's rows, as views of its plan: a forward - that runs one row apart reads and writes that row's slot and no other.""" - pool = build_pool(max_slots=6) - rids = ["a", "b", "c"] - for rid in rids: - pool.ingest_request(rid) - step = decode_step(rids) - ctx = StepContext(request_ids=rids, graph_walk="decode", slot=0, capture=False) - assert pool.admit(step, ctx).ok - rows = pool.plan(step, ctx)["main"] - one = rows.narrow(1) - assert one.num_rows == 1 and one.slot_indices.tolist() == rows.slot_indices[1:2].tolist() - assert list(one.host_slot_indices) == list(rows.host_slot_indices[1:2]) - before = pool.block("conv", 0).clone() - pool.scatter_("conv", one, torch.ones(1, 96, 3)) - changed = [i for i in range(6) if not torch.equal(pool.block("conv", 0)[i], before[i])] - assert changed == [pool._slots["b"]["main"].index] - torch.testing.assert_close(pool.gather("conv", rows)[1:2], pool.gather("conv", one)) From aaa25e23619128082833d38740d17031c7ec8e05 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 02:17:15 -0700 Subject: [PATCH 34/69] nemotron_duplex tests: capture on CPU as main's runner does; the engine fake has a capture_runner PiecewiseAcceleratorGraphRunner.warmup_and_capture takes the planned buckets since main's capture budget. The CPU capture helper passes None, as the engine does on a device with no plan: every bucket is warmed up and captured. That warm-up now reads the accelerator's memory stats, which a CPU device has none of, so the helper stubs them. The batch cap test's submodule entry carries capture_runner, which Engine.get_max_batch_size now reads. Co-Authored-By: Claude Opus 5.5 --- test/modular/test_nemotron_duplex_submodules.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index 5a98e0f20d..8e0cd1a7ca 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -413,9 +413,14 @@ def replay(self): device=torch.device("cpu"), autocast_dtype=None, num_slots=1, node_name=label, memory_pool=AcceleratorGraphPool("cpu"), device_capability=SimpleNamespace(graph_backend=backend), ) - with mock.patch.object(cgr, "capture_into_graph", lambda fn, *a: (_Replay(fn), fn())): + # 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() + runner.warmup_and_capture(None) assert runner.any_graphs and not runner.dropped_shapes runner.calls = 0 run = runner.run @@ -892,7 +897,8 @@ def test_encoder_and_codec_take_no_batch_wider_than_their_graphs(): "audio_codec": AudioCodecDecoderSubmodule(codec=nn.Identity(), config=NemotronDuplexConfig()), } engine = Engine.__new__(Engine) - engine._submodules = {name: SimpleNamespace(submodule=node, accelerator_graph_runner=None) + 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 From e4030eefb56b86afc84cb43e90ad5ca6cdeb44da Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 12:04:20 -0700 Subject: [PATCH 35/69] test_recurrent_padding_rows: drop the stale select wording from a docstring Co-Authored-By: Claude Opus 5.5 --- test/modular/test_recurrent_padding_rows.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/modular/test_recurrent_padding_rows.py b/test/modular/test_recurrent_padding_rows.py index cd9887b923..c401befe81 100644 --- a/test/modular/test_recurrent_padding_rows.py +++ b/test/modular/test_recurrent_padding_rows.py @@ -150,8 +150,8 @@ def test_a_real_request_padding_and_release_across_the_rotation(): def test_gather_scatter_follow_the_step_addressing(): """``gather`` / ``scatter_`` address the step's rows: a replay's padding rows - read and write the sink and nothing else, and a ``select`` of the step - addresses the same slots as the step.""" + read and write the sink and nothing else, and a step planned over a subset + of the requests addresses the same slots as the full step.""" pool = build_pool(max_slots=6) step = decode_step(["a", "b"]) eager = StepContext(request_ids=["a", "b"], graph_walk="decode", slot=0, capture=False) From 791c31f648319cd7b2c3bec70be6d17a7f6b824e Mon Sep 17 00:00:00 2001 From: Naomi Sagan Date: Wed, 21 Oct 2026 16:02:14 -0700 Subject: [PATCH 36/69] Bounded KV: a sink + window per-request KV resource Extracted unchanged from #383 (model/minicpm-speech), the state its commits 61a92cc0f..04fa7b249 leave the bounded-KV package in: the resource (mstar/engine/resources/kv/bounded/), its tests (test/modular/test_bounded_kv.py) and its sections of docs/adding_models.rst. Bounded KV only: token2wav and the rest of #383 stay on that PR, which rebases onto this. #383 adds no bounded-KV exports to mstar/engine/resources/__init__.py, so neither does this. Co-Authored-By: Claude Opus 5.5 --- docs/adding_models.rst | 38 +++ mstar/engine/resources/kv/bounded/__init__.py | 9 + mstar/engine/resources/kv/bounded/config.py | 131 +++++++++ mstar/engine/resources/kv/bounded/kernels.py | 245 ++++++++++++++++ mstar/engine/resources/kv/bounded/layout.py | 86 ++++++ mstar/engine/resources/kv/bounded/manager.py | 272 ++++++++++++++++++ test/modular/test_bounded_kv.py | 258 +++++++++++++++++ 7 files changed, 1039 insertions(+) create mode 100644 mstar/engine/resources/kv/bounded/__init__.py create mode 100644 mstar/engine/resources/kv/bounded/config.py create mode 100644 mstar/engine/resources/kv/bounded/kernels.py create mode 100644 mstar/engine/resources/kv/bounded/layout.py create mode 100644 mstar/engine/resources/kv/bounded/manager.py create mode 100644 test/modular/test_bounded_kv.py diff --git a/docs/adding_models.rst b/docs/adding_models.rst index b33d9cc57a..b4ce9774b8 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -393,6 +393,16 @@ The spec types are: - A fixed-capacity frame ring. It adds ``tokens_per_frame``, one ``RingKVLayerConfig`` per layer, and the deployment-tunable ``num_sessions``. Ring storage is currently paired with FlexAttention. + * - ``BoundedKVSpec(config=BoundedKVConfig(...))`` + - Fixed-size per-request KV with sink + window retention, and its attention. A + request's stream is an optional read-only source (e.g. a voice prompt's keys, + shared by the requests that use it) followed by what it writes; ``retention`` maps + the source length to the ``SinkWindow`` kept (the first ``sink`` entries and the + last ``window``). Each step declares its rows' ``StreamPosition`` and writes + (``BoundedKVStep``), so rows at different positions batch, and nothing is + reserved or released per step. ``attend`` takes an optional per-distance score + term (``rel_bias``) for relative-position attention. Float32 by default (Triton + kernels). See `Bounded (sink + window) caches`_. * - ``AttentionSpec(config=AttentionConfig(kv_cache=...))`` - Self-attention planned over the named cache. ``backend`` selects ``AttnBackend.AUTO`` (the default, which lets the engine pick the platform's @@ -675,6 +685,32 @@ LTX-2.5 is the reference (``mstar/model/ltx2_5/submodules.py``). us to zero-pad to 128. So that tower attends through the resource only when it is replaying a captured graph, and runs flash-attn varlen otherwise. +Bounded (sink + window) caches +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +A streaming decoder that keeps a fixed amount of context per request (a prompt or voice +prefix plus the last N steps, as StreamingLLM-style sinks or a streaming vocoder's chunk +caches do) declares a ``BoundedKVSpec``. A request takes a slot when it first runs and +keeps it; the window is a ring overwritten in place, so a step reserves and releases +nothing. A step's layout is at most five key ranges and three write ranges per row, +computed on the host from the position the step declares, which is also what lets rows at +different positions share a batch and a captured graph. + +The model calls ``attend(layer_idx, q, k, v, source, plan)``: each row attends its +retained keys in stream order and its own, unmasked, and its writes land after. Pass +``source`` (``[sources, rows, H, L, 2 * head_dim]``, keys then values) when requests start +from shared prefixes that the slots should not copy; each row reads the one its +``StreamPosition.source`` names, so rows of different prefixes share a batch and a graph. ``rel_bias`` adds a term that depends only on +the query-key distance, for Transformer-XL-style relative positions. MiniCPM-o's token2wav +(``CACHE_FAMILIES`` in ``mstar/model/minicpm_o/components/token2wav.py``) declares three: +its two conformer caches and its DiT's. + +Two flags cover retention rules that are not quite a token FIFO. +``SinkWindow(window_includes_step=True)`` counts the step's own tokens toward the window. +``BoundedKVConfig(reverse_step_order=True)`` appends each step's tokens last-first, which +only makes sense for attention that ignores key order; it decides which of the oldest +step's tokens leave first. + Step 3 — Declare the computation graph -------------------------------------- @@ -1416,6 +1452,8 @@ that a misspelled setting is never silently ignored: ``prefix_cache``, ``prefix_cache_salt`` * - ``KVSpec`` with ``RingKVConfig`` - ``num_sessions`` + * - ``BoundedKVSpec`` + - ``max_slots`` * - ``AttentionSpec`` - ``backend`` (``auto`` / ``flashinfer`` / ``dense`` / ``flex`` / ``xpu_paged``), ``flashinfer_backend`` (``auto`` / ``fa2`` / ``fa3``) diff --git a/mstar/engine/resources/kv/bounded/__init__.py b/mstar/engine/resources/kv/bounded/__init__.py new file mode 100644 index 0000000000..591920f96a --- /dev/null +++ b/mstar/engine/resources/kv/bounded/__init__.py @@ -0,0 +1,9 @@ +from mstar.engine.resources.kv.bounded.config import ( + BoundedKVConfig, + BoundedKVSpec, + BoundedKVStep, + SinkWindow, + StreamPosition, +) + +__all__ = ["BoundedKVConfig", "BoundedKVSpec", "BoundedKVStep", "SinkWindow", "StreamPosition"] diff --git a/mstar/engine/resources/kv/bounded/config.py b/mstar/engine/resources/kv/bounded/config.py new file mode 100644 index 0000000000..b72cd6eb5c --- /dev/null +++ b/mstar/engine/resources/kv/bounded/config.py @@ -0,0 +1,131 @@ +"""What a model declares about a bounded KV cache. + +A request's cache is a logical stream: an optional read-only source prefix +(a voice prompt's keys, say, shared by every request that uses it) followed by +the tokens the request writes. Retention keeps the stream's first ``sink`` +entries and its last ``window`` (StreamingLLM's attention sinks plus a sliding +window), so a request's footprint is fixed and known when it starts: each slot +holds only the written part of the sink and a ring for the window, and source +entries are read where they live. +""" + +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, NamedTuple + +import torch + +from mstar.engine.resources.spec import NodeResourceSpec +from mstar.engine.resources.step import ResourceStep + +if TYPE_CHECKING: + from mstar.engine.resources.base import Resource + + +class SinkWindow(NamedTuple): + """Keep the stream's first ``sink`` and last ``window`` entries. With + ``window_includes_step`` the window counts the step's own tokens, so a step + that writes ``span`` attends only the last ``window - span`` retained ones.""" + + sink: int + window: int + window_includes_step: bool = False + + +@dataclass +class BoundedKVConfig: + num_layers: int + num_heads: int + head_dim: int + # key/value rows a request owns per layer, e.g. 2 for batched CFG + rows_per_request: int + # the longest source prefix a request may have; sizes the slots + max_source_len: int + # source length -> what that request keeps + retention: Callable[[int], SinkWindow] + max_slots: int = 32 + # the shortest source prefix a request may have: the slots also hold the part + # of a sink its source doesn't cover, which is largest for the shortest source + min_source_len: int = 0 + dtype: torch.dtype = torch.float32 + # A step's tokens join the stream last-first. Only for attention that + # ignores key order, where it decides which of the oldest step's tokens are + # evicted first: Step-Audio2's flow decoder keeps the front of each chunk. + reverse_step_order: bool = False + + def __post_init__(self): + if self.max_slots < 1: + raise ValueError(f"max_slots must be >= 1, got {self.max_slots}") + if not 0 <= self.min_source_len <= self.max_source_len: + raise ValueError(f"source lengths {self.min_source_len}..{self.max_source_len} are not a range") + sink = window = retained = 0 + for n in range(self.min_source_len, self.max_source_len + 1): + policy = self.retention(n) + if policy.sink < 0 or policy.window < 0: + raise ValueError(f"retention for {n} source tokens is {policy}") + sink = max(sink, policy.sink - n) + window = max(window, policy.window) + retained = max(retained, policy.sink + policy.window) + # fixed once here: a plan reads them for every row + self._sink_capacity, self._window_capacity, self._max_retained = sink, window, retained + + @property + def sink_capacity(self) -> int: + """Written tokens a slot's sink region holds (the sink past the source).""" + return self._sink_capacity + + @property + def window_capacity(self) -> int: + return self._window_capacity + + @property + def max_retained(self) -> int: + """The most entries a request retains, over every allowed source length.""" + return self._max_retained + + @property + def slot_tokens(self) -> int: + return self.sink_capacity + self.window_capacity + + @property + def slot_bytes(self) -> int: + per_token = self.rows_per_request * self.num_heads * 2 * self.head_dim + return self.num_layers * self.slot_tokens * per_token * torch.empty((), dtype=self.dtype).element_size() + + +@dataclass +class BoundedKVSpec(NodeResourceSpec): + config: BoundedKVConfig + + @property + def resource_class(self) -> "type[Resource]": + from mstar.engine.resources.kv.bounded.manager import BoundedKVManager + + return BoundedKVManager + + def apply_yaml_overrides(self, max_slots: int | None = None): + if max_slots is not None: + self.config.max_slots = max_slots + + +class StreamPosition(NamedTuple): + """Where a request's stream stands when a step starts.""" + + source_len: int + # tokens the request has written before this step + written: int + # which of the sources ``attend`` is given this request's starts with; -1 for its + # own entry of the ``row_source`` it is given + source: int = 0 + + +@dataclass(frozen=True) +class BoundedKVStep(ResourceStep): + """Each segment's ``span`` is what its request writes this step (0 reads only). + + The positions come from the model, not the resource, so a plan never waits + on the previous step's commit and planning the same step twice (a piecewise + region inside a forward that declared it too) is harmless. + """ + + positions: Mapping[str, StreamPosition] = field(default_factory=dict) diff --git a/mstar/engine/resources/kv/bounded/kernels.py b/mstar/engine/resources/kv/bounded/kernels.py new file mode 100644 index 0000000000..9c6f78b79a --- /dev/null +++ b/mstar/engine/resources/kv/bounded/kernels.py @@ -0,0 +1,245 @@ +"""Attention over a bounded KV step's ranges, and its write-back (Triton). + +Keys and values share a ``2 * head_dim`` last axis in the slot and the source: +keys in ``[0, D)``, values in ``[D, 2D)``. Every row attends to its retained +ranges and to this step's own tokens, unmasked; dot products use TF32 unless +``ieee``. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from mstar.engine.resources.kv.bounded.layout import NUM_READS, NUM_WRITES, ROW_HEAD, ROW_INTS + + +@triton.jit +def _scores(q, k, bias_rows, rel, valid, sm_scale, HAS_BIAS: tl.constexpr, PRECISION: tl.constexpr): + """``q . k`` plus, with a bias, each pair's entry ``rel`` of its query's band row.""" + s = tl.dot(q, tl.trans(k), input_precision=PRECISION) + if HAS_BIAS: + s += tl.load(bias_rows + rel, mask=valid, other=0.0) + return tl.where(valid, s * sm_scale, float("-inf")) + + +@triton.jit +def _bounded_attention_kernel( + q_ptr, k_ptr, v_ptr, o_ptr, cache_ptr, source_ptr, row_source_ptr, row_slots_ptr, table_ptr, bias_ptr, + T, H, sm_scale, + s_qn, s_qh, s_qt, + s_cs, s_cr, s_ch, s_cp, + s_sv, s_sr, s_sh, s_sp, + s_rv, s_rr, s_rh, s_rp, + s_on, s_ot, s_oh, + s_bn, s_bh, s_bt, + ROWS: tl.constexpr, HAS_SOURCE: tl.constexpr, HAS_ROW_SOURCE: tl.constexpr, HAS_BIAS: tl.constexpr, + ROW_INTS: tl.constexpr, HEAD: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, D: tl.constexpr, PRECISION: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_nh = tl.program_id(1) + n = pid_nh // H + h = pid_nh % H + b = n // ROWS + c = n % ROWS + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, D) + mask_m = offs_m < T + q = tl.load(q_ptr + n * s_qn + h * s_qh + offs_m[:, None] * s_qt + offs_d[None, :], + mask=mask_m[:, None], other=0.0) + bias_rows = bias_ptr + n * s_bn + h * s_bh + offs_m[:, None] * s_bt + + row = table_ptr + b * ROW_INTS + slot = tl.load(row).to(tl.int64) + slot_base = cache_ptr + slot * s_cs + c * s_cr + h * s_ch + # a row's source: one of the shared ones by index, or (-1) its own slot of the row sources + src = tl.load(row + 1).to(tl.int64) + source_base = source_ptr + src * s_sv + c * s_sr + h * s_sh + s_pos = s_sp + if HAS_ROW_SOURCE: + if src < 0: + own = tl.load(row_slots_ptr + b).to(tl.int64) + source_base = row_source_ptr + own * s_rv + c * s_rr + h * s_rh + s_pos = s_rp + # The five ranges (source, slot, source, slot, slot: layout.READ_FROM_SOURCE), + # in stream order: key j of them is at stream position j, this step's own at + # ``total + t``. Tiled as one sequence, so short ranges do not each pay a + # masked tile. + s0 = tl.load(row + HEAD + 0) + e1 = tl.load(row + HEAD + 1) + s1 = tl.load(row + HEAD + 2) + e2 = e1 + tl.load(row + HEAD + 3) + s2 = tl.load(row + HEAD + 4) + e3 = e2 + tl.load(row + HEAD + 5) + s3 = tl.load(row + HEAD + 6) + e4 = e3 + tl.load(row + HEAD + 7) + s4 = tl.load(row + HEAD + 8) + total = e4 + tl.load(row + HEAD + 9) + + m_i = tl.full([BLOCK_M], float("-inf"), tl.float32) + l_i = tl.zeros([BLOCK_M], tl.float32) + acc = tl.zeros([BLOCK_M, D], tl.float32) + + # this step's own tokens (separate k and v tensors, laid out like q); a pair's + # band entry is its query position minus its key position, plus T - 1 + for j0 in range(0, T, BLOCK_N): + offs_n = j0 + tl.arange(0, BLOCK_N) + mask_n = offs_n < T + kv_off = n * s_qn + h * s_qh + offs_n[:, None] * s_qt + offs_d[None, :] + k = tl.load(k_ptr + kv_off, mask=mask_n[:, None], other=0.0) + v = tl.load(v_ptr + kv_off, mask=mask_n[:, None], other=0.0) + rel = offs_m[:, None] - offs_n[None, :] + T - 1 + s = _scores(q, k, bias_rows, rel, mask_m[:, None] & mask_n[None, :], sm_scale, HAS_BIAS, PRECISION) + m_new = tl.maximum(m_i, tl.max(s, axis=1)) + alpha = tl.exp(m_i - m_new) + p = tl.exp(s - m_new[:, None]) + l_i = l_i * alpha + tl.sum(p, axis=1) + acc = acc * alpha[:, None] + tl.dot(p, v, input_precision=PRECISION) + m_i = m_new + + for j0 in range(0, total, BLOCK_N): + j = j0 + tl.arange(0, BLOCK_N) + valid = j < total + in_source = (j < e1) | ((j >= e2) & (j < e3)) + pos = tl.where(j < e1, s0 + j, + tl.where(j < e2, s1 + j - e1, + tl.where(j < e3, s2 + j - e2, + tl.where(j < e4, s3 + j - e3, s4 + j - e4)))) + ptrs = slot_base + pos[:, None] * s_cp + offs_d[None, :] + if HAS_SOURCE: + source_ptrs = source_base + pos[:, None] * s_pos + offs_d[None, :] + ptrs = tl.where(in_source[:, None], source_ptrs, ptrs) + k = tl.load(ptrs, mask=valid[:, None], other=0.0) + v = tl.load(ptrs + D, mask=valid[:, None], other=0.0) + rel = total + offs_m[:, None] - j[None, :] + T - 1 + s = _scores(q, k, bias_rows, rel, mask_m[:, None] & valid[None, :], sm_scale, HAS_BIAS, PRECISION) + m_new = tl.maximum(m_i, tl.max(s, axis=1)) + alpha = tl.exp(m_i - m_new) + p = tl.exp(s - m_new[:, None]) + l_i = l_i * alpha + tl.sum(p, axis=1) + acc = acc * alpha[:, None] + tl.dot(p, v, input_precision=PRECISION) + m_i = m_new + + out = acc / l_i[:, None] + tl.store(o_ptr + n * s_on + offs_m[:, None] * s_ot + h * s_oh + offs_d[None, :], out, + mask=mask_m[:, None]) + + +@triton.jit +def _bounded_store_kernel( + k_ptr, v_ptr, cache_ptr, table_ptr, + H, + s_kn, s_kh, s_kt, + s_cs, s_cr, s_ch, s_cp, + ROWS: tl.constexpr, ROW_INTS: tl.constexpr, NUM_READS: tl.constexpr, REVERSE: tl.constexpr, + HEAD: tl.constexpr, + BLOCK_T: tl.constexpr, D: tl.constexpr, +): + pid_t = tl.program_id(0) + pid_nh = tl.program_id(1) + n = pid_nh // H + h = pid_nh % H + b = n // ROWS + c = n % ROWS + row = table_ptr + b * ROW_INTS + slot = tl.load(row).to(tl.int64) + offs_t = pid_t * BLOCK_T + tl.arange(0, BLOCK_T) + offs_d = tl.arange(0, D) + for w in tl.static_range(3): + base = row + HEAD + 2 * NUM_READS + 3 * w + offset = tl.load(base) + start = tl.load(base + 1) + count = tl.load(base + 2) + mask = offs_t < count + fresh = offset - offs_t if REVERSE else offset + offs_t + src = n * s_kn + h * s_kh + fresh[:, None] * s_kt + offs_d[None, :] + k = tl.load(k_ptr + src, mask=mask[:, None]) + v = tl.load(v_ptr + src, mask=mask[:, None]) + dst = cache_ptr + slot * s_cs + c * s_cr + h * s_ch + (start + offs_t)[:, None] * s_cp + offs_d[None, :] + tl.store(dst, k, mask=mask[:, None]) + tl.store(dst + D, v, mask=mask[:, None]) + + +def bounded_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cache: torch.Tensor, + source: torch.Tensor | None, + table: torch.Tensor, + rel_bias: torch.Tensor | None = None, + ieee: bool = False, + row_source: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> torch.Tensor: + """``q, k, v [B * rows, H, T, D]`` (contiguous) over their retained ranges: + ``cache [slots, rows, H, cap, 2D]`` and ``source [sources, rows, H, L, 2D]`` (each + row's own, by the table's index) per ``table [>= B, ROW_INTS]``. A row whose source + index is -1 reads ``row_source = (storage [slots, rows, H, L', 2D], slots [>= B])`` + at its slot instead. ``rel_bias [B * rows, H, T, R]`` adds to the + score of query ``i`` and stream key ``j`` its entry ``(n + i) - j + T - 1``, + ``n`` being the row's retained count (Transformer-XL's position term, which + depends on the distance alone). Returns ``[B * rows, T, H, D]``.""" + n, h, t, d = q.shape + rows = cache.shape[1] + assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() + assert cache.stride(-1) == 1 and table.shape[-1] == ROW_INTS + has_row_source = row_source is not None + has_source = source is not None or has_row_source + if source is None: + source = cache + s_src = (0, 0, 0, 0) + else: + assert source.dim() == 5 and source.stride(-1) == 1 + s_src = (source.stride(0), source.stride(1), source.stride(2), source.stride(3)) + if has_row_source: + row_storage, row_slots = row_source + assert row_storage.dim() == 5 and row_storage.stride(-1) == 1 + s_row = (row_storage.stride(0), row_storage.stride(1), row_storage.stride(2), row_storage.stride(3)) + else: + row_storage, row_slots, s_row = cache, table, (0, 0, 0, 0) + has_bias = rel_bias is not None + if has_bias: + assert rel_bias.shape[:3] == (n, h, t) and rel_bias.stride(-1) == 1 + s_bias = (rel_bias.stride(0), rel_bias.stride(1), rel_bias.stride(2)) + else: + rel_bias, s_bias = q, (0, 0, 0) + out = torch.empty(n, t, h, d, device=q.device, dtype=q.dtype) + block_m = 16 if n * h < 128 else 64 + grid = (triton.cdiv(t, block_m), n * h) + _bounded_attention_kernel[grid]( + q, k, v, out, cache, source, row_storage, row_slots, table, rel_bias, + t, h, d ** -0.5, + q.stride(0), q.stride(1), q.stride(2), + cache.stride(0), cache.stride(1), cache.stride(2), cache.stride(3), + *s_src, + *s_row, + out.stride(0), out.stride(1), out.stride(2), + *s_bias, + ROWS=rows, HAS_SOURCE=has_source, HAS_ROW_SOURCE=has_row_source, HAS_BIAS=has_bias, + ROW_INTS=ROW_INTS, HEAD=ROW_HEAD, + BLOCK_M=block_m, BLOCK_N=64, D=d, PRECISION="ieee" if ieee else "tf32", + ) + return out + + +def bounded_store( + k: torch.Tensor, v: torch.Tensor, cache: torch.Tensor, table: torch.Tensor, reverse: bool = False, +) -> None: + """This step's ``k, v [B * rows, H, T, D]`` into the slots per ``table``'s writes + (a range's tokens counting down from its offset when ``reverse``).""" + n, h, t, d = k.shape + assert k.is_contiguous() and v.is_contiguous() and cache.stride(-1) == 1 + block_t = 64 + grid = (triton.cdiv(t, block_t), n * h) + _bounded_store_kernel[grid]( + k, v, cache, table, h, + k.stride(0), k.stride(1), k.stride(2), + cache.stride(0), cache.stride(1), cache.stride(2), cache.stride(3), + ROWS=cache.shape[1], ROW_INTS=ROW_INTS, NUM_READS=NUM_READS, REVERSE=reverse, HEAD=ROW_HEAD, + BLOCK_T=block_t, D=d, + ) + + +assert NUM_READS == 5 and NUM_WRITES == 3, "the kernels unroll five reads and three writes" diff --git a/mstar/engine/resources/kv/bounded/layout.py b/mstar/engine/resources/kv/bounded/layout.py new file mode 100644 index 0000000000..5bff64a103 --- /dev/null +++ b/mstar/engine/resources/kv/bounded/layout.py @@ -0,0 +1,86 @@ +"""Host-side layout of one row of a bounded KV step. + +A row's retained keys are at most five contiguous ranges, in stream order: +the sink's source part, the sink's written part, the window's source part and +the window's ring (two ranges when it wraps). Its writes are at most three: +the sink's unfilled part and the ring (again two when it wraps). Each is a +``(start, length)`` pair in the source or the slot, so a step's layout is a +small integer table, and the kernels need nothing else. +""" + +from typing import NamedTuple + +from mstar.engine.resources.kv.bounded.config import SinkWindow + +# reads: (start, len) pairs in this order; READ_FROM_SOURCE marks which index +# the source rather than the slot +NUM_READS = 5 +READ_FROM_SOURCE = (True, False, True, False, False) +# writes: (fresh offset of the range's first entry, slot start, len) triples +NUM_WRITES = 3 +# slot, source, reads, writes +ROW_HEAD = 2 +ROW_INTS = ROW_HEAD + 2 * NUM_READS + 3 * NUM_WRITES + + +class RowLayout(NamedTuple): + reads: tuple[tuple[int, int], ...] + writes: tuple[tuple[int, int, int], ...] + + +def _ring_ranges(first: int, count: int, window: int) -> list[tuple[int, int]]: + """``count`` ring positions from ``first`` (mod ``window``) as at most two ranges.""" + if count <= 0: + return [] + start = first % window + head = min(count, window - start) + return [(start, head)] + ([(0, count - head)] if count > head else []) + + +def row_layout( + source_len: int, written: int, span: int, policy: SinkWindow, sink_capacity: int, + reverse: bool = False, +) -> RowLayout: + """The reads and writes of a row whose stream holds ``source_len + written`` + entries and gains ``span`` this step. Slot offsets put the sink's written + part at ``[0, sink_capacity)`` and the ring after it. With ``reverse`` a + step's tokens join the stream last-first (see ``BoundedKVConfig``).""" + sink, window, includes_step = policy + n = source_len + written + ring = sink_capacity + + reads: list[tuple[int, int]] = [(0, min(sink, source_len))] + reads.append((0, max(0, min(sink, n) - source_len))) + lo = max(sink, n - (window - span if includes_step else window)) + reads.append((lo, max(0, source_len - lo))) + first = max(lo, source_len) + wrapped = _ring_ranges(first - sink, n - first, window) + reads += [(ring + s, c) for s, c in wrapped] + [(0, 0)] * (2 - len(wrapped)) + + def fresh(i: int) -> int: + """Which of this step's tokens is stream entry ``i``.""" + return n + span - 1 - i if reverse else i - n + + writes: list[tuple[int, int, int]] = [] + sink_end = min(sink, n + span) + if sink_end > n: + writes.append((fresh(n), n - source_len, sink_end - n)) + if window > 0: + first = max(sink, n, n + span - window) + for s, c in _ring_ranges(first - sink, n + span - first, window): + writes.append((fresh(first), ring + s, c)) + first += c + writes += [(0, 0, 0)] * (NUM_WRITES - len(writes)) + return RowLayout(tuple(reads), tuple(writes)) + + +EMPTY_ROW = RowLayout(((0, 0),) * NUM_READS, ((0, 0, 0),) * NUM_WRITES) + + +def flatten_row(slot: int, layout: RowLayout, source: int = 0) -> list[int]: + out = [slot, source] + for start, count in layout.reads: + out += (start, count) + for offset, start, count in layout.writes: + out += (offset, start, count) + return out diff --git a/mstar/engine/resources/kv/bounded/manager.py b/mstar/engine/resources/kv/bounded/manager.py new file mode 100644 index 0000000000..35e1cd443a --- /dev/null +++ b/mstar/engine/resources/kv/bounded/manager.py @@ -0,0 +1,272 @@ +"""A pool of fixed-size per-request KV slots with sink + window retention. + +Unlike the paged ``KVManager`` there is nothing to reserve per step and +nothing to release: a request takes a slot when it first runs and keeps it, +the window is a ring overwritten in place, and a step's layout is a small +integer table computed on the host from positions the step declares. +""" + +from __future__ import annotations + +import threading +from typing import NamedTuple + +import torch + +from mstar.engine.resources.base import CGSlotKey, CGSlotSpec, EngineResourceInfo, Resource +from mstar.engine.resources.kv.bounded.config import BoundedKVConfig, BoundedKVSpec, BoundedKVStep +from mstar.engine.resources.kv.bounded.layout import ( + EMPTY_ROW, + NUM_READS, + READ_FROM_SOURCE, + ROW_HEAD, + ROW_INTS, + RowLayout, + flatten_row, + row_layout, +) +from mstar.engine.resources.spec import ResourceReqConfig +from mstar.engine.resources.step import ADMIT_OK, AdmitOutcome, AdmitRuntimeError, AllocationFailed, StepContext +from mstar.utils.h2d import PinnedStager + + +class BoundedPlan(NamedTuple): + table: torch.Tensor # [>= rows, ROW_INTS] int32 + # (slot, layout, source) per row, padding included; what the table holds + rows: tuple[tuple[int, RowLayout, int], ...] + + def retained(self) -> torch.Tensor: + """Each row's retained keys before its step, ``[rows]`` int32 on the device.""" + return self.table[: len(self.rows), ROW_HEAD + 1:ROW_HEAD + 2 * NUM_READS:2].sum(1) + + + +class _Request: + __slots__ = ("slot", "written") + + def __init__(self, slot: int): + self.slot = slot + self.written = 0 + + +class BoundedKVManager(Resource): + @classmethod + def build(cls, spec: BoundedKVSpec, info: EngineResourceInfo) -> "BoundedKVManager": + if info.joint_comm_group is not None and info.joint_comm_group.world_size > 1: + raise NotImplementedError("bounded KV does not shard across ranks yet") + return cls(spec.config, info.device) + + def __init__(self, config: BoundedKVConfig, device: torch.device): + self.config = config + self._device = torch.device(device) + c = config + # layer-major so one layer's slots are a contiguous [slots, rows, H, cap, 2D] + self._cache = torch.zeros( + c.num_layers, c.max_slots, c.rows_per_request, c.num_heads, max(c.slot_tokens, 1), + 2 * c.head_dim, dtype=c.dtype, device=self._device, + ) + self._requests: dict[str, _Request] = {} + self._free = list(reversed(range(c.max_slots))) + self._lock = threading.RLock() + + self._stager = PinnedStager(torch.int32) + self._cg_max_bs = 0 + self._cg_tables: dict[CGSlotKey, torch.Tensor] = {} + self._eager_table: torch.Tensor | None = None + self._current: BoundedPlan | None = None + + # Model-facing + + def layer(self, layer_idx: int) -> torch.Tensor: + """One layer's slots, ``[max_slots, rows, H, cap, 2D]``.""" + return self._cache[layer_idx] + + @property + def current(self) -> BoundedPlan: + if self._current is None: + raise RuntimeError("bounded KV has no plan for this step; declare a BoundedKVStep") + return self._current + + def attend( + self, + layer_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + source: torch.Tensor | None = None, + plan: BoundedPlan | None = None, + rel_bias: torch.Tensor | None = None, + ieee: bool = False, + row_source: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + """Attention for ``q, k, v [B * rows, H, T, D]`` over each row's retained + keys (``source [sources, rows, H, L, 2D]`` is this layer's source prefixes, + each row's picked by its ``StreamPosition.source``) and this + step's own, then this step's keys into the slots. ``[B * rows, T, H, D]``. + ``plan`` defaults to the step's. ``rel_bias`` is a per-distance score term + (``kernels.bounded_attention``), for relative-position attention. Rows whose + source is -1 read theirs from ``row_source = (storage [slots, rows, H, L, 2D], + slots [>= B])``, e.g. another resource's per-request entries.""" + plan = plan or self.current + cache = self._cache[layer_idx] + if not q.is_cuda: + return self._attend_torch(q, k, v, cache, source, plan, rel_bias, row_source) + from mstar.engine.resources.kv.bounded.kernels import bounded_attention, bounded_store + + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + out = bounded_attention(q, k, v, cache, source, plan.table, rel_bias, ieee, row_source) + bounded_store(k, v, cache, plan.table, self.config.reverse_step_order) + return out + + def _attend_torch( + self, q, k, v, cache, source, plan: BoundedPlan, rel_bias=None, row_source=None, + ) -> torch.Tensor: + """The kernels' semantics row by row, from the host layouts (CPU).""" + n, h, t, d = q.shape + rows = self.config.rows_per_request + out = q.new_empty(n, t, h, d) + own = row_source[1].tolist() if row_source is not None else None + for i in range(n): + b, c = divmod(i, rows) + slot, layout, src = plan.rows[b] + row_src = row_source[0][own[b], c] if src < 0 else (None if source is None else source[src, c]) + keys, values = [], [] + for from_source, (start, count) in zip(READ_FROM_SOURCE, layout.reads, strict=True): + if count == 0 or (from_source and row_src is None): + continue + kv = (row_src if from_source else cache[slot, c])[:, start:start + count] + keys.append(kv[..., :d]) + values.append(kv[..., d:]) + keys, values = torch.cat([*keys, k[i]], 1), torch.cat([*values, v[i]], 1) + scores = q[i] @ keys.transpose(-1, -2) + if rel_bias is not None: + total = keys.shape[1] - t + # query i sits at total + i, key j at j + rel = total + torch.arange(t)[:, None] - torch.arange(keys.shape[1])[None, :] + t - 1 + scores = scores + rel_bias[i].gather(-1, rel.to(q.device).expand(h, -1, -1)) + att = torch.softmax(scores * d ** -0.5, dim=-1) + out[i] = (att @ values).transpose(0, 1) + for i in range(n): + b, c = divmod(i, rows) + slot, layout, _ = plan.rows[b] + for offset, start, count in layout.writes: + if count == 0: + continue + if self.config.reverse_step_order: + fresh = torch.arange(offset, offset - count, -1, device=k.device) + else: + fresh = torch.arange(offset, offset + count, device=k.device) + cache[slot, c, :, start:start + count] = torch.cat([k[i][:, fresh], v[i][:, fresh]], dim=-1) + return out + + # Request lifecycle + + def ingest_request(self, rid: str, overrides: ResourceReqConfig | None = None): + del overrides + + def remove_request(self, rid: str): + self.reset_request(rid, free=True) + + def reset_request(self, rid: str, free: bool = False): + """Nothing to zero: a slot's contents are only read where a later + position says they were written.""" + with self._lock: + req = self._requests.get(rid) + if req is None: + return + if free: + del self._requests[rid] + self._free.append(req.slot) + else: + req.written = 0 + + # Step lifecycle + + def admit(self, step: BoundedKVStep, ctx: StepContext) -> AdmitOutcome: + if ctx.capture: + return ADMIT_OK + with self._lock: + for seg in step.segments or (): + if ctx.is_padding_row(seg.request_id): + continue + pos = step.positions.get(seg.request_id) + if pos is None: + return AdmitOutcome(ok=False, ready=False, reason=AdmitRuntimeError( + f"bounded KV step declares no position for {seg.request_id!r}")) + if not self.config.min_source_len <= pos.source_len <= self.config.max_source_len: + return AdmitOutcome(ok=False, ready=False, reason=AdmitRuntimeError( + f"{seg.request_id!r} has a {pos.source_len}-token source; the slots " + f"take {self.config.min_source_len} to {self.config.max_source_len}")) + req = self._requests.get(seg.request_id) + if req is None: + if not self._free: + return AdmitOutcome(ok=False, reason=AllocationFailed( + message=f"bounded KV is full: {self.config.max_slots} slots", + pages_short=1, label=seg.label, request_id=seg.request_id)) + req = self._requests[seg.request_id] = _Request(self._free.pop()) + if pos.written != req.written: + return AdmitOutcome(ok=False, ready=False, reason=AdmitRuntimeError( + f"{seg.request_id!r} declares {pos.written} written tokens; " + f"{req.written} were committed")) + return ADMIT_OK + + def plan(self, step: BoundedKVStep, ctx: StepContext) -> BoundedPlan: + c = self.config + values: list[int] = [] + rows = [] + for seg in step.segments or (): + req = None if ctx.capture or ctx.is_padding_row(seg.request_id) else self._requests.get(seg.request_id) + if req is None: + # padding rows read and write nothing + rows.append((0, EMPTY_ROW, 0)) + else: + pos = step.positions[seg.request_id] + rows.append((req.slot, row_layout(pos.source_len, pos.written, seg.span, + c.retention(pos.source_len), c.sink_capacity, + c.reverse_step_order), pos.source)) + values += flatten_row(*rows[-1]) + table = self._table(ctx, len(rows)) + self._stager.copy_(table, values, pad_value=0) + self._current = BoundedPlan(table, tuple(rows)) + return self._current + + def commit(self, step: BoundedKVStep, ctx: StepContext) -> None: + if ctx.capture: + return + with self._lock: + for seg in step.segments or (): + req = self._requests.get(seg.request_id) + pos = step.positions.get(seg.request_id) + if req is not None and pos is not None and not ctx.is_padding_row(seg.request_id): + # from the declared position, so committing a step twice is harmless + req.written = pos.written + seg.span + + @property + def force_double_buffer(self): + return True + + def _table(self, ctx: StepContext, rows: int) -> torch.Tensor: + lease = ctx.slot_lease + if lease is None: + if self._eager_table is None or self._eager_table.shape[0] < rows: + self._eager_table = torch.zeros(max(rows, 1), ROW_INTS, dtype=torch.int32, device=self._device) + return self._eager_table[:max(rows, 1)] + key = CGSlotKey(bucket=lease.bucket, slot=lease.slot, label="main") + table = self._cg_tables.get(key) + if table is None: + size = max(self._cg_max_bs, lease.bucket.bs, rows) + table = self._cg_tables[key] = torch.zeros(size, ROW_INTS, dtype=torch.int32, device=self._device) + return table + + def build_cuda_graph_buffers(self, slots: list[CGSlotSpec], max_bs: int, max_seq_len: int) -> None: + del slots, max_seq_len + self._cg_max_bs = max(self._cg_max_bs, max_bs) + + def post_warmup_validate(self): + if self._requests: + raise RuntimeError(f"bounded KV slots still held after capture: {sorted(self._requests)}") + + def cleanup(self): + self._cg_tables.clear() + self._eager_table = None + self._current = None diff --git a/test/modular/test_bounded_kv.py b/test/modular/test_bounded_kv.py new file mode 100644 index 0000000000..d33c280973 --- /dev/null +++ b/test/modular/test_bounded_kv.py @@ -0,0 +1,258 @@ +"""Bounded (sink + window) KV: the host layout against a brute-force stream, +token2wav's three cache rules expressed as sink + window, and the kernels +against torch.""" + +import random + +import pytest +import torch + +from mstar.engine.resources.kv.bounded.config import SinkWindow +from mstar.engine.resources.kv.bounded.layout import ( + READ_FROM_SOURCE, + flatten_row, + row_layout, +) + + +class SlotSim: + """Applies layouts to a list-valued slot and reads back token ids.""" + + def __init__(self, source: list, policy: SinkWindow, sink_capacity: int, reverse: bool = False): + self.reverse = reverse + self.source = source + self.policy = policy + self.sink_capacity = sink_capacity + self.slot = [None] * (sink_capacity + policy.window) + self.written = 0 + + def step(self, fresh: list, write: bool = True) -> list: + """The ids a step attends to (retained, then its own), then its writes.""" + span = len(fresh) if write else 0 + layout = row_layout(len(self.source), self.written, span, self.policy, self.sink_capacity, + self.reverse) + keys = [] + for from_source, (start, count) in zip(READ_FROM_SOURCE, layout.reads, strict=True): + buf = self.source if from_source else self.slot + keys += buf[start:start + count] + assert None not in keys, "read an unwritten slot entry" + for offset, start, count in layout.writes: + if self.reverse: + self.slot[start:start + count] = fresh[offset - count + 1:offset + 1][::-1] + else: + self.slot[start:start + count] = fresh[offset:offset + count] + self.written += span + return keys + fresh + + +@pytest.mark.parametrize("reverse", [False, True]) +@pytest.mark.parametrize("seed", range(20)) +def test_layout_matches_stream(seed, reverse): + rng = random.Random(seed) + source_len = rng.randrange(0, 40) + policy = SinkWindow(rng.randrange(0, 50), rng.randrange(0, 30)) + sink_capacity = max(0, policy.sink - source_len) + source = [("s", i) for i in range(source_len)] + sim = SlotSim(source, policy, sink_capacity, reverse) + stream = list(source) + for step in range(12): + fresh = [("g", len(stream) - source_len + j) for j in range(rng.randrange(1, 20))] + n = len(stream) + retained = stream[:min(policy.sink, n)] + stream[max(policy.sink, n - policy.window):n] + assert sim.step(fresh) == retained + fresh, step + stream += fresh[::-1] if reverse else fresh + + +@pytest.mark.parametrize("source_len", [0, 5, 29, 30, 60]) +def test_slots_hold_the_sink_of_the_shortest_source(source_len): + """A fixed sink longer than a short source spills into the slot; the slots are + sized for the shortest source allowed, so a request with a longer one (whose + retention alone needs no sink region) shares the layout.""" + from mstar.engine.resources.kv.bounded.config import BoundedKVConfig + + config = BoundedKVConfig( + num_layers=1, num_heads=1, head_dim=1, rows_per_request=1, + min_source_len=0, max_source_len=60, retention=lambda n: SinkWindow(30, 20), + ) + assert config.sink_capacity == 30 and config.max_retained == 50 + source = [("s", i) for i in range(source_len)] + sim = SlotSim(source, config.retention(source_len), config.sink_capacity) + stream = list(source) + for step in range(8): + fresh = [("g", len(stream) - source_len + j) for j in range(7)] + n = len(stream) + assert sim.step(fresh) == stream[:min(30, n)] + stream[max(30, n - 20):n] + fresh, step + stream += fresh + + +# token2wav: a voice of P tokens (2P frames at 50 Hz); 28-code windows advance by 25 +P, KEEP = 60, 100 + + +def _upstream(windows: int): + """Upstream's caches as id lists, window by window (``Token2WavState.truncate``): + returns what each window attends to, per cache.""" + p2, limit = 2 * P, 2 * P + KEEP + voice1 = [("v1", i) for i in range(P)] + voice2 = [("v2", i) for i in range(p2)] + voiced = [("vd", i) for i in range(p2)] + enc1, enc2, dit = list(voice1), list(voice2), list(voiced) + seen = [] + for w in range(windows): + new1 = [("g1", 25 * w + j) for j in range(25)] + new2 = [("g2", 50 * w + j) for j in range(50)] + newd = [("gd", 50 * w + j) for j in range(50)] + seen.append({"enc1": enc1 + new1, "enc2": enc2 + new2, "dit": newd + dit}) + enc1, enc2, dit = enc1 + new1, enc2 + new2, newd + dit + if len(dit) > limit: + dit = dit[:p2] + dit[len(dit) - KEEP:] + if len(enc2) > limit: + n2, n1 = len(enc2), len(enc1) + kept = list(range(p2)) + list(range(n2 - KEEP, n2)) + enc1 = [enc1[i % n1] for i in kept[:limit // 2]] + enc2 = enc2[:p2] + enc2[n2 - KEEP:] + return seen + + +def test_token2wav_caches_are_sink_window(): + p2 = 2 * P + windows = 9 + upstream = _upstream(windows) + voice1 = [("v1", i) for i in range(P)] + voice2 = [("v2", i) for i in range(p2)] + voiced = [("vd", i) for i in range(p2)] + # The DiT's attention is order-free, and upstream keeps the first 2P of [newest + # window, ..., oldest, voice] (each window in its own order): sink + window + # over a stream of the voice's tail, then the rest of it and each window + # last-first. + dit_source = voiced[p2 - KEEP:] + voiced[:p2 - KEEP][::-1] + caches = { + "enc1": SlotSim(voice1, SinkWindow(P + KEEP // 2, 0), KEEP // 2), + "enc2": SlotSim(voice2, SinkWindow(p2, KEEP), 0), + "dit": SlotSim(dit_source, SinkWindow(KEEP, p2), 0, reverse=True), + } + for w in range(windows): + fresh = { + "enc1": [("g1", 25 * w + j) for j in range(25)], + "enc2": [("g2", 50 * w + j) for j in range(50)], + "dit": [("gd", 50 * w + j) for j in range(50)], + } + for name, sim in caches.items(): + got = sim.step(fresh[name]) + want = upstream[w][name] + if name == "dit": + assert sorted(got) == sorted(want), (w, name) + else: + # the conformer's attention is positional: same order + assert got == want, (w, name) + + +def _reference(q, k, v, cache, source, rows_layout, slots, sources, rel_bias=None, row_source=None): + """Torch: attention over each row's ranges (in stream order) and its own keys, + plus each pair's position term; then the writes.""" + n, h, t, d = q.shape + rows = cache.shape[1] + out = torch.empty(n, t, h, d, dtype=q.dtype, device=q.device) + new_cache = cache.clone() + for i in range(n): + b, c = divmod(i, rows) + layout, slot = rows_layout[b], slots[b] + parts = [] + for from_source, (start, count) in zip(READ_FROM_SOURCE, layout.reads, strict=True): + if not from_source: + buf = cache[slot, c] + elif sources[b] < 0: + buf = row_source[0][row_source[1][b], c] + else: + buf = source[sources[b], c] + parts.append(buf[:, start:start + count]) + kv = torch.cat(parts, dim=1) + total = kv.shape[1] + keys = torch.cat([kv[..., :d], k[i]], dim=1) + values = torch.cat([kv[..., d:], v[i]], dim=1) + scores = q[i] @ keys.transpose(-1, -2) + if rel_bias is not None: + for qi in range(t): + for j in range(total + t): + scores[:, qi, j] += rel_bias[i, :, qi, total + qi - j + t - 1] + att = torch.softmax(scores * d ** -0.5, dim=-1) + out[i] = (att @ values).transpose(0, 1) + for offset, start, count in layout.writes: + new_cache[slot, c, :, start:start + count, :d] = k[i][:, offset:offset + count] + new_cache[slot, c, :, start:start + count, d:] = v[i][:, offset:offset + count] + return out, new_cache + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernels need CUDA") +@pytest.mark.parametrize("with_row_source", [False, True]) +@pytest.mark.parametrize("with_bias", [False, True]) +@pytest.mark.parametrize("with_source", [True, False]) +def test_kernels_match_torch(with_source, with_bias, with_row_source): + from mstar.engine.resources.kv.bounded.kernels import bounded_attention, bounded_store + + torch.manual_seed(0) + rng = random.Random(0) + dev = "cuda" + rows, h, d, t = 2, 4, 64, 50 + source_len = 120 if with_source else 0 + policy = SinkWindow(30 if with_source else 0, 90) + sink_capacity = max(0, policy.sink - source_len) + slots_n = 6 + cache = torch.randn(slots_n, rows, h, sink_capacity + policy.window, 2 * d, device=dev) + # three sources (voices, say), each row starting from its own + source = torch.randn(3, rows, h, max(source_len, 1), 2 * d, device=dev) if with_source else None + # rows at different positions, including a read-only one + written = [0, 50, 100, 170, 260] + spans = [t, t, t, t, 0] + slots = [3, 0, 5, 1, 2] + sources = [2, 0, 1, 2, 0] if with_source else [0] * 5 + row_source = None + if with_row_source and with_source: + # rows 1 and 3 read their own entry of a second, per-row source + sources = [2, -1, 1, -1, 0] + row_source = (torch.randn(4, rows, h, source_len + 7, 2 * d, device=dev), + torch.tensor([0, 3, 0, 1, 0], dtype=torch.int32, device=dev)) + layouts = [row_layout(source_len, w, s, policy, sink_capacity) for w, s in zip(written, spans, strict=True)] + table = torch.tensor([flatten_row(s, lay, src) for s, lay, src in zip(slots, layouts, sources, strict=True)], + dtype=torch.int32, device=dev) + b = len(slots) + q, k, v = (torch.randn(b * rows, h, t, d, device=dev) for _ in range(3)) + src = source if with_source else torch.zeros(1, rows, h, 1, 2 * d, device=dev) + longest = max(sum(c for _, c in lay.reads) for lay in layouts) + bias = torch.randn(b * rows, h, t, longest + 2 * t - 1, device=dev) * 4 if with_bias else None + # float64, so the reference is not itself TF32 + ref_rows = None if row_source is None else (row_source[0].double(), row_source[1].tolist()) + want, want_cache = _reference(q.double(), k.double(), v.double(), cache.double(), src.double(), + layouts, slots, sources, None if bias is None else bias.double(), ref_rows) + want, want_cache = want.float(), want_cache.float() + got = bounded_attention(q, k, v, cache, source, table, bias, ieee=True, row_source=row_source) + bounded_store(k, v, cache, table) + torch.testing.assert_close(got, want, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(cache, want_cache, rtol=0, atol=0) + + +@pytest.mark.parametrize("ring_buckets,dilation", [(4, 1), (6, 1), (3, 4), (5, 8)]) +def test_waypoint_ring_is_sink_window(ring_buckets, dilation): + """Waypoint's ring (``RingKVManager`` + flex visibility) as a sink + window: no + sink, a window of ``ring_buckets`` frames that counts the current one, over a + stream that only every ``dilation``-th frame writes to (the others attend + without writing). Compared slot set by slot set against its block tables.""" + from types import SimpleNamespace + + from mstar.engine.resources.attn.flex import FlexAttentionManager + + tokens = 128 + manager = SimpleNamespace(_kv_config=SimpleNamespace(tokens_per_frame=tokens)) + geometry = (ring_buckets, ring_buckets, dilation) + policy = SinkWindow(0, ring_buckets * tokens, window_includes_step=True) + written = 0 + for frame in range(4 * ring_buckets * dilation): + visible = FlexAttentionManager._visible_blocks_for(manager, geometry, session_idx=0, frame_pos=frame) + # their ring slots, without the scratch frame the current one sits in + want = sorted({b // (tokens // 128) for b in visible} - {ring_buckets}) + span = tokens if frame % dilation == 0 else 0 + layout = row_layout(0, written, span, policy, 0) + got = sorted({pos // tokens for start, count in layout.reads[3:] for pos in range(start, start + count)}) + assert got == want, (frame, got, want) + assert sum(c for _, c in layout.reads[:3]) == 0 + written += span From e2cd2e1d9bd312d804dc93b13d40d6307987d0b4 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:07:18 -0700 Subject: [PATCH 37/69] Bounded KV: a per-query causal window SinkWindow(right=...) makes the window per query: each query attends the keys at distances [-right, window] from it, so a step of several tokens is causal within itself and each of them sees exactly `window` back, as a streaming conformer's att_context_size does. The default stays per step (every query sees the retained keys and all of the step's own), which is token2wav's behaviour. Only right=0 is accepted, with no sink and the stream in order, and the window must be the same for every source length, since the kernel takes it as one scalar (BoundedKVConfig.query_window). The kernel masks each query to distances [0, window], bounds its tile loops to [i - window, i] for the block's queries, and indexes rel_bias by distance (R = window + 1); the torch path does the same. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/kv/bounded/config.py | 27 +++++++++- mstar/engine/resources/kv/bounded/kernels.py | 57 +++++++++++++++----- mstar/engine/resources/kv/bounded/layout.py | 2 +- mstar/engine/resources/kv/bounded/manager.py | 23 +++++--- 4 files changed, 86 insertions(+), 23 deletions(-) diff --git a/mstar/engine/resources/kv/bounded/config.py b/mstar/engine/resources/kv/bounded/config.py index b72cd6eb5c..e320684c6a 100644 --- a/mstar/engine/resources/kv/bounded/config.py +++ b/mstar/engine/resources/kv/bounded/config.py @@ -25,11 +25,19 @@ class SinkWindow(NamedTuple): """Keep the stream's first ``sink`` and last ``window`` entries. With ``window_includes_step`` the window counts the step's own tokens, so a step - that writes ``span`` attends only the last ``window - span`` retained ones.""" + that writes ``span`` attends only the last ``window - span`` retained ones. + + By default the window is per step: every query of a step attends the same + keys, the retained ones and all of the step's own. With ``right`` set it is + per query: each query attends the keys at distances ``[-right, window]`` + from it, so a step of several tokens is causal within itself and each of + them sees exactly ``window`` back (a streaming conformer's + ``att_context_size``). Only ``right=0``, with no sink, is supported so far.""" sink: int window: int window_includes_step: bool = False + right: int | None = None @dataclass @@ -59,15 +67,27 @@ def __post_init__(self): if not 0 <= self.min_source_len <= self.max_source_len: raise ValueError(f"source lengths {self.min_source_len}..{self.max_source_len} are not a range") sink = window = retained = 0 + query_windows = set() for n in range(self.min_source_len, self.max_source_len + 1): policy = self.retention(n) if policy.sink < 0 or policy.window < 0: raise ValueError(f"retention for {n} source tokens is {policy}") + if policy.right is not None and ( + policy.right != 0 or policy.sink or policy.window_includes_step or self.reverse_step_order + ): + raise ValueError(f"retention for {n} source tokens is {policy}: a per-query window " + "takes right=0, no sink, and the stream in order") + query_windows.add(None if policy.right is None else policy.window) sink = max(sink, policy.sink - n) window = max(window, policy.window) retained = max(retained, policy.sink + policy.window) + # the kernel takes a per-query window as one scalar for the whole step + if len(query_windows) > 1: + raise ValueError(f"retention mixes windows {sorted(query_windows, key=str)}; a per-query " + "window must be the same for every source length") # fixed once here: a plan reads them for every row self._sink_capacity, self._window_capacity, self._max_retained = sink, window, retained + self._query_window = next(iter(query_windows)) @property def sink_capacity(self) -> int: @@ -78,6 +98,11 @@ def sink_capacity(self) -> int: def window_capacity(self) -> int: return self._window_capacity + @property + def query_window(self) -> int | None: + """Each query's window when retention is per query (``SinkWindow.right``), else None.""" + return self._query_window + @property def max_retained(self) -> int: """The most entries a request retains, over every allowed source length.""" diff --git a/mstar/engine/resources/kv/bounded/kernels.py b/mstar/engine/resources/kv/bounded/kernels.py index 9c6f78b79a..57eeabcc5b 100644 --- a/mstar/engine/resources/kv/bounded/kernels.py +++ b/mstar/engine/resources/kv/bounded/kernels.py @@ -2,8 +2,9 @@ Keys and values share a ``2 * head_dim`` last axis in the slot and the source: keys in ``[0, D)``, values in ``[D, 2D)``. Every row attends to its retained -ranges and to this step's own tokens, unmasked; dot products use TF32 unless -``ieee``. +ranges and to this step's own tokens: unmasked, or with a per-query window +(``SinkWindow.right``) each query to the keys at distances ``[0, window]``. +Dot products use TF32 unless ``ieee``. """ from __future__ import annotations @@ -27,7 +28,7 @@ def _scores(q, k, bias_rows, rel, valid, sm_scale, HAS_BIAS: tl.constexpr, PRECI @triton.jit def _bounded_attention_kernel( q_ptr, k_ptr, v_ptr, o_ptr, cache_ptr, source_ptr, row_source_ptr, row_slots_ptr, table_ptr, bias_ptr, - T, H, sm_scale, + T, H, sm_scale, window, s_qn, s_qh, s_qt, s_cs, s_cr, s_ch, s_cp, s_sv, s_sr, s_sh, s_sp, @@ -35,7 +36,7 @@ def _bounded_attention_kernel( s_on, s_ot, s_oh, s_bn, s_bh, s_bt, ROWS: tl.constexpr, HAS_SOURCE: tl.constexpr, HAS_ROW_SOURCE: tl.constexpr, HAS_BIAS: tl.constexpr, - ROW_INTS: tl.constexpr, HEAD: tl.constexpr, + PER_QUERY: tl.constexpr, ROW_INTS: tl.constexpr, HEAD: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, D: tl.constexpr, PRECISION: tl.constexpr, ): pid_m = tl.program_id(0) @@ -82,16 +83,34 @@ def _bounded_attention_kernel( l_i = tl.zeros([BLOCK_M], tl.float32) acc = tl.zeros([BLOCK_M, D], tl.float32) + # A per-query window bounds each loop to the keys some query of this block + # sees: query i (stream position total + i) sees back to i - window. + if PER_QUERY: + m_lo = pid_m * BLOCK_M + own_lo = tl.maximum(m_lo - window, 0) + own_hi = tl.minimum(m_lo + BLOCK_M, T) + old_lo = tl.maximum(total + m_lo - window, 0) + else: + own_lo = 0 + own_hi = T + old_lo = 0 + # this step's own tokens (separate k and v tensors, laid out like q); a pair's - # band entry is its query position minus its key position, plus T - 1 - for j0 in range(0, T, BLOCK_N): + # band entry is its query position minus its key position, plus T - 1, or + # with a per-query window the distance alone + for j0 in range(own_lo, own_hi, BLOCK_N): offs_n = j0 + tl.arange(0, BLOCK_N) mask_n = offs_n < T kv_off = n * s_qn + h * s_qh + offs_n[:, None] * s_qt + offs_d[None, :] k = tl.load(k_ptr + kv_off, mask=mask_n[:, None], other=0.0) v = tl.load(v_ptr + kv_off, mask=mask_n[:, None], other=0.0) - rel = offs_m[:, None] - offs_n[None, :] + T - 1 - s = _scores(q, k, bias_rows, rel, mask_m[:, None] & mask_n[None, :], sm_scale, HAS_BIAS, PRECISION) + valid = mask_m[:, None] & mask_n[None, :] + if PER_QUERY: + rel = offs_m[:, None] - offs_n[None, :] + valid = valid & (rel >= 0) & (rel <= window) + else: + rel = offs_m[:, None] - offs_n[None, :] + T - 1 + s = _scores(q, k, bias_rows, rel, valid, sm_scale, HAS_BIAS, PRECISION) m_new = tl.maximum(m_i, tl.max(s, axis=1)) alpha = tl.exp(m_i - m_new) p = tl.exp(s - m_new[:, None]) @@ -99,7 +118,7 @@ def _bounded_attention_kernel( acc = acc * alpha[:, None] + tl.dot(p, v, input_precision=PRECISION) m_i = m_new - for j0 in range(0, total, BLOCK_N): + for j0 in range(old_lo, total, BLOCK_N): j = j0 + tl.arange(0, BLOCK_N) valid = j < total in_source = (j < e1) | ((j >= e2) & (j < e3)) @@ -113,8 +132,13 @@ def _bounded_attention_kernel( ptrs = tl.where(in_source[:, None], source_ptrs, ptrs) k = tl.load(ptrs, mask=valid[:, None], other=0.0) v = tl.load(ptrs + D, mask=valid[:, None], other=0.0) - rel = total + offs_m[:, None] - j[None, :] + T - 1 - s = _scores(q, k, bias_rows, rel, mask_m[:, None] & valid[None, :], sm_scale, HAS_BIAS, PRECISION) + seen = mask_m[:, None] & valid[None, :] + if PER_QUERY: + rel = total + offs_m[:, None] - j[None, :] + seen = seen & (rel <= window) + else: + rel = total + offs_m[:, None] - j[None, :] + T - 1 + s = _scores(q, k, bias_rows, rel, seen, sm_scale, HAS_BIAS, PRECISION) m_new = tl.maximum(m_i, tl.max(s, axis=1)) alpha = tl.exp(m_i - m_new) p = tl.exp(s - m_new[:, None]) @@ -172,6 +196,7 @@ def bounded_attention( rel_bias: torch.Tensor | None = None, ieee: bool = False, row_source: tuple[torch.Tensor, torch.Tensor] | None = None, + window: int | None = None, ) -> torch.Tensor: """``q, k, v [B * rows, H, T, D]`` (contiguous) over their retained ranges: ``cache [slots, rows, H, cap, 2D]`` and ``source [sources, rows, H, L, 2D]`` (each @@ -180,7 +205,10 @@ def bounded_attention( at its slot instead. ``rel_bias [B * rows, H, T, R]`` adds to the score of query ``i`` and stream key ``j`` its entry ``(n + i) - j + T - 1``, ``n`` being the row's retained count (Transformer-XL's position term, which - depends on the distance alone). Returns ``[B * rows, T, H, D]``.""" + depends on the distance alone). With a per-query ``window`` + (``BoundedKVConfig.query_window``) query ``i`` attends only the keys at + distances ``[0, window]`` from it, and the band is indexed by that distance, + ``R = window + 1``. Returns ``[B * rows, T, H, D]``.""" n, h, t, d = q.shape rows = cache.shape[1] assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() @@ -202,6 +230,7 @@ def bounded_attention( has_bias = rel_bias is not None if has_bias: assert rel_bias.shape[:3] == (n, h, t) and rel_bias.stride(-1) == 1 + assert window is None or rel_bias.shape[3] > window, "a per-query band holds window + 1 distances" s_bias = (rel_bias.stride(0), rel_bias.stride(1), rel_bias.stride(2)) else: rel_bias, s_bias = q, (0, 0, 0) @@ -210,7 +239,7 @@ def bounded_attention( grid = (triton.cdiv(t, block_m), n * h) _bounded_attention_kernel[grid]( q, k, v, out, cache, source, row_storage, row_slots, table, rel_bias, - t, h, d ** -0.5, + t, h, d ** -0.5, 0 if window is None else window, q.stride(0), q.stride(1), q.stride(2), cache.stride(0), cache.stride(1), cache.stride(2), cache.stride(3), *s_src, @@ -218,7 +247,7 @@ def bounded_attention( out.stride(0), out.stride(1), out.stride(2), *s_bias, ROWS=rows, HAS_SOURCE=has_source, HAS_ROW_SOURCE=has_row_source, HAS_BIAS=has_bias, - ROW_INTS=ROW_INTS, HEAD=ROW_HEAD, + PER_QUERY=window is not None, ROW_INTS=ROW_INTS, HEAD=ROW_HEAD, BLOCK_M=block_m, BLOCK_N=64, D=d, PRECISION="ieee" if ieee else "tf32", ) return out diff --git a/mstar/engine/resources/kv/bounded/layout.py b/mstar/engine/resources/kv/bounded/layout.py index 5bff64a103..4c511f46a1 100644 --- a/mstar/engine/resources/kv/bounded/layout.py +++ b/mstar/engine/resources/kv/bounded/layout.py @@ -45,7 +45,7 @@ def row_layout( entries and gains ``span`` this step. Slot offsets put the sink's written part at ``[0, sink_capacity)`` and the ring after it. With ``reverse`` a step's tokens join the stream last-first (see ``BoundedKVConfig``).""" - sink, window, includes_step = policy + sink, window, includes_step = policy.sink, policy.window, policy.window_includes_step n = source_len + written ring = sink_capacity diff --git a/mstar/engine/resources/kv/bounded/manager.py b/mstar/engine/resources/kv/bounded/manager.py index 35e1cd443a..5ba590d974 100644 --- a/mstar/engine/resources/kv/bounded/manager.py +++ b/mstar/engine/resources/kv/bounded/manager.py @@ -104,7 +104,9 @@ def attend( each row's picked by its ``StreamPosition.source``) and this step's own, then this step's keys into the slots. ``[B * rows, T, H, D]``. ``plan`` defaults to the step's. ``rel_bias`` is a per-distance score term - (``kernels.bounded_attention``), for relative-position attention. Rows whose + (``kernels.bounded_attention``), for relative-position attention; with a + per-query window (``SinkWindow.right``) it is indexed by distance, + ``[B * rows, H, T, window + 1]``. Rows whose source is -1 read theirs from ``row_source = (storage [slots, rows, H, L, 2D], slots [>= B])``, e.g. another resource's per-request entries.""" plan = plan or self.current @@ -114,7 +116,8 @@ def attend( from mstar.engine.resources.kv.bounded.kernels import bounded_attention, bounded_store q, k, v = q.contiguous(), k.contiguous(), v.contiguous() - out = bounded_attention(q, k, v, cache, source, plan.table, rel_bias, ieee, row_source) + out = bounded_attention(q, k, v, cache, source, plan.table, rel_bias, ieee, row_source, + self.config.query_window) bounded_store(k, v, cache, plan.table, self.config.reverse_step_order) return out @@ -124,6 +127,7 @@ def _attend_torch( """The kernels' semantics row by row, from the host layouts (CPU).""" n, h, t, d = q.shape rows = self.config.rows_per_request + window = self.config.query_window out = q.new_empty(n, t, h, d) own = row_source[1].tolist() if row_source is not None else None for i in range(n): @@ -139,11 +143,16 @@ def _attend_torch( values.append(kv[..., d:]) keys, values = torch.cat([*keys, k[i]], 1), torch.cat([*values, v[i]], 1) scores = q[i] @ keys.transpose(-1, -2) - if rel_bias is not None: - total = keys.shape[1] - t - # query i sits at total + i, key j at j - rel = total + torch.arange(t)[:, None] - torch.arange(keys.shape[1])[None, :] + t - 1 - scores = scores + rel_bias[i].gather(-1, rel.to(q.device).expand(h, -1, -1)) + # query i sits at total + i, key j at j + total = keys.shape[1] - t + dist = (total + torch.arange(t)[:, None] - torch.arange(keys.shape[1])[None, :]).to(q.device) + if window is None: + if rel_bias is not None: + scores = scores + rel_bias[i].gather(-1, (dist + t - 1).expand(h, -1, -1)) + else: + if rel_bias is not None: + scores = scores + rel_bias[i].gather(-1, dist.clamp(0, window).expand(h, -1, -1)) + scores = scores.masked_fill((dist < 0) | (dist > window), float("-inf")) att = torch.softmax(scores * d ** -0.5, dim=-1) out[i] = (att @ values).transpose(0, 1) for i in range(n): From 951c94a2f33aeec43b9588c2b7154391e0602103 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:07:29 -0700 Subject: [PATCH 38/69] Bounded KV: guard the online softmax against a fully masked tile A query whose keys so far are all masked has a running max of -inf, and the rescale exp(m_i - m_new) is then exp(-inf + inf) = NaN, which poisons its accumulator for good. With a per-query window this is the ordinary case: a step longer than the window has key tiles that the block's later queries do not see at all. Offset such rows by 0, so alpha and p come out 0. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/kv/bounded/kernels.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/mstar/engine/resources/kv/bounded/kernels.py b/mstar/engine/resources/kv/bounded/kernels.py index 57eeabcc5b..b63ce463a9 100644 --- a/mstar/engine/resources/kv/bounded/kernels.py +++ b/mstar/engine/resources/kv/bounded/kernels.py @@ -112,8 +112,11 @@ def _bounded_attention_kernel( rel = offs_m[:, None] - offs_n[None, :] + T - 1 s = _scores(q, k, bias_rows, rel, valid, sm_scale, HAS_BIAS, PRECISION) m_new = tl.maximum(m_i, tl.max(s, axis=1)) - alpha = tl.exp(m_i - m_new) - p = tl.exp(s - m_new[:, None]) + # a query whose keys so far are all masked keeps m = -inf, and + # exp(-inf - -inf) is NaN: offset by 0 instead, so alpha and p are 0 + m_off = tl.where(m_new == float("-inf"), 0.0, m_new) + alpha = tl.exp(m_i - m_off) + p = tl.exp(s - m_off[:, None]) l_i = l_i * alpha + tl.sum(p, axis=1) acc = acc * alpha[:, None] + tl.dot(p, v, input_precision=PRECISION) m_i = m_new @@ -140,8 +143,9 @@ def _bounded_attention_kernel( rel = total + offs_m[:, None] - j[None, :] + T - 1 s = _scores(q, k, bias_rows, rel, seen, sm_scale, HAS_BIAS, PRECISION) m_new = tl.maximum(m_i, tl.max(s, axis=1)) - alpha = tl.exp(m_i - m_new) - p = tl.exp(s - m_new[:, None]) + m_off = tl.where(m_new == float("-inf"), 0.0, m_new) + alpha = tl.exp(m_i - m_off) + p = tl.exp(s - m_off[:, None]) l_i = l_i * alpha + tl.sum(p, axis=1) acc = acc * alpha[:, None] + tl.dot(p, v, input_precision=PRECISION) m_i = m_new From 7c9d2cba4eed269e2b50f847032dd1df0738c8be Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:07:29 -0700 Subject: [PATCH 39/69] Bounded KV: cast the probabilities to the values' dtype before the dot tl.dot takes both operands in one dtype, and p is float32, so a bf16 cache has float32 against bf16 in the second dot. Cast p to the values' dtype, as FlashAttention-style kernels do. Not run here: the bf16 kernel test added with the other tests exercises it. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/kv/bounded/kernels.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mstar/engine/resources/kv/bounded/kernels.py b/mstar/engine/resources/kv/bounded/kernels.py index b63ce463a9..02b3c874b2 100644 --- a/mstar/engine/resources/kv/bounded/kernels.py +++ b/mstar/engine/resources/kv/bounded/kernels.py @@ -118,7 +118,7 @@ def _bounded_attention_kernel( alpha = tl.exp(m_i - m_off) p = tl.exp(s - m_off[:, None]) l_i = l_i * alpha + tl.sum(p, axis=1) - acc = acc * alpha[:, None] + tl.dot(p, v, input_precision=PRECISION) + acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v, input_precision=PRECISION) m_i = m_new for j0 in range(old_lo, total, BLOCK_N): @@ -147,7 +147,7 @@ def _bounded_attention_kernel( alpha = tl.exp(m_i - m_off) p = tl.exp(s - m_off[:, None]) l_i = l_i * alpha + tl.sum(p, axis=1) - acc = acc * alpha[:, None] + tl.dot(p, v, input_precision=PRECISION) + acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v, input_precision=PRECISION) m_i = m_new out = acc / l_i[:, None] From a51fc80df68b8c75aefab9a3c613563120a810b2 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:07:43 -0700 Subject: [PATCH 40/69] Bounded KV: dtype=None is the engine's KV dtype Settled in build from EngineResourceInfo.kv_dtype, so a model whose layers compute in the engine's dtype keeps its cache in that dtype without naming it. The default stays float32, so token2wav's caches are unchanged. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/kv/bounded/config.py | 4 +++- mstar/engine/resources/kv/bounded/manager.py | 2 ++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/mstar/engine/resources/kv/bounded/config.py b/mstar/engine/resources/kv/bounded/config.py index e320684c6a..8b83000976 100644 --- a/mstar/engine/resources/kv/bounded/config.py +++ b/mstar/engine/resources/kv/bounded/config.py @@ -55,7 +55,9 @@ class BoundedKVConfig: # the shortest source prefix a request may have: the slots also hold the part # of a sink its source doesn't cover, which is largest for the shortest source min_source_len: int = 0 - dtype: torch.dtype = torch.float32 + # None is the engine's KV dtype (``EngineResourceInfo.kv_dtype``), settled when + # the resource is built + dtype: torch.dtype | None = torch.float32 # A step's tokens join the stream last-first. Only for attention that # ignores key order, where it decides which of the oldest step's tokens are # evicted first: Step-Audio2's flow decoder keeps the front of each chunk. diff --git a/mstar/engine/resources/kv/bounded/manager.py b/mstar/engine/resources/kv/bounded/manager.py index 5ba590d974..42379ca8ba 100644 --- a/mstar/engine/resources/kv/bounded/manager.py +++ b/mstar/engine/resources/kv/bounded/manager.py @@ -54,6 +54,8 @@ class BoundedKVManager(Resource): def build(cls, spec: BoundedKVSpec, info: EngineResourceInfo) -> "BoundedKVManager": if info.joint_comm_group is not None and info.joint_comm_group.world_size > 1: raise NotImplementedError("bounded KV does not shard across ranks yet") + if spec.config.dtype is None: + spec.config.dtype = info.kv_dtype if info.kv_dtype is not None else torch.get_default_dtype() return cls(spec.config, info.device) def __init__(self, config: BoundedKVConfig, device: torch.device): From 69c7a7a8d3b2dd21253545652a89d20322b8cc3c Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:07:50 -0700 Subject: [PATCH 41/69] Bounded KV: rename the graph-buffer hook to main's name Main calls build_accelerator_graph_buffers (same signature); under the old name, build_cuda_graph_buffers, the override never ran. No behaviour changes today, since each captured table is also sized by its own lease's bucket, but the hook now records the widest capture as intended. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/kv/bounded/manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mstar/engine/resources/kv/bounded/manager.py b/mstar/engine/resources/kv/bounded/manager.py index 42379ca8ba..826ae7537a 100644 --- a/mstar/engine/resources/kv/bounded/manager.py +++ b/mstar/engine/resources/kv/bounded/manager.py @@ -269,7 +269,7 @@ def _table(self, ctx: StepContext, rows: int) -> torch.Tensor: table = self._cg_tables[key] = torch.zeros(size, ROW_INTS, dtype=torch.int32, device=self._device) return table - def build_cuda_graph_buffers(self, slots: list[CGSlotSpec], max_bs: int, max_seq_len: int) -> None: + def build_accelerator_graph_buffers(self, slots: list[CGSlotSpec], max_bs: int, max_seq_len: int) -> None: del slots, max_seq_len self._cg_max_bs = max(self._cg_max_bs, max_bs) From 8cdc9e8f24381b983a359a31dd52d361d19d38c3 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:08:01 -0700 Subject: [PATCH 42/69] Bounded KV: admit refusals name the request at fault AdmitRuntimeError with no request_ids is a refusal of the step as a whole, so one call's position mismatch failed every request in its batch. Name it, as the ring KV manager does, so only that request fails. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/kv/bounded/manager.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/mstar/engine/resources/kv/bounded/manager.py b/mstar/engine/resources/kv/bounded/manager.py index 826ae7537a..4297d0eec6 100644 --- a/mstar/engine/resources/kv/bounded/manager.py +++ b/mstar/engine/resources/kv/bounded/manager.py @@ -203,11 +203,13 @@ def admit(self, step: BoundedKVStep, ctx: StepContext) -> AdmitOutcome: pos = step.positions.get(seg.request_id) if pos is None: return AdmitOutcome(ok=False, ready=False, reason=AdmitRuntimeError( - f"bounded KV step declares no position for {seg.request_id!r}")) + f"bounded KV step declares no position for {seg.request_id!r}", + request_ids=frozenset({seg.request_id}))) if not self.config.min_source_len <= pos.source_len <= self.config.max_source_len: return AdmitOutcome(ok=False, ready=False, reason=AdmitRuntimeError( f"{seg.request_id!r} has a {pos.source_len}-token source; the slots " - f"take {self.config.min_source_len} to {self.config.max_source_len}")) + f"take {self.config.min_source_len} to {self.config.max_source_len}", + request_ids=frozenset({seg.request_id}))) req = self._requests.get(seg.request_id) if req is None: if not self._free: @@ -218,7 +220,7 @@ def admit(self, step: BoundedKVStep, ctx: StepContext) -> AdmitOutcome: if pos.written != req.written: return AdmitOutcome(ok=False, ready=False, reason=AdmitRuntimeError( f"{seg.request_id!r} declares {pos.written} written tokens; " - f"{req.written} were committed")) + f"{req.written} were committed", request_ids=frozenset({seg.request_id}))) return ADMIT_OK def plan(self, step: BoundedKVStep, ctx: StepContext) -> BoundedPlan: From 32fb8fc00b2f5fbf4d44135833c1c8489208f7a8 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:08:19 -0700 Subject: [PATCH 43/69] Bounded KV: commit skips span-0 segments Commit sets a row's count from its declared position, so a step that declares a row with span 0 would put the count back to that position. When a forward's own step reads a row and a piecewise region inside it writes the row, the region commits first and the outer step's commit then undid it. A span-0 segment writes nothing, so it has nothing to commit. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/kv/bounded/manager.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mstar/engine/resources/kv/bounded/manager.py b/mstar/engine/resources/kv/bounded/manager.py index 4297d0eec6..e1abb47426 100644 --- a/mstar/engine/resources/kv/bounded/manager.py +++ b/mstar/engine/resources/kv/bounded/manager.py @@ -248,6 +248,10 @@ def commit(self, step: BoundedKVStep, ctx: StepContext) -> None: return with self._lock: for seg in step.segments or (): + if seg.span == 0: + # writes nothing; skipped so a step that only reads a row + # (an outer step whose region writes it) can't undo the write + continue req = self._requests.get(seg.request_id) pos = step.positions.get(seg.request_id) if req is not None and pos is not None and not ctx.is_padding_row(seg.request_id): From 8cff7bbd9d6e95bc717158df588fa1e6df2dfbe9 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:08:41 -0700 Subject: [PATCH 44/69] Bounded KV: drop force_double_buffer, whose hazard this resource lacks The flag guards plan(N+1) overwriting a staging buffer or a graph-read buffer before step N has consumed it. Here the host side is a PinnedStager ring, which waits on a buffer's copy event before rewriting it, and the resource does not pre-plan (supports_preplan is False), so StepRunner.plan runs it only on the GPU thread, on the stream that then runs the step: plan(N+1)'s H2D into a table is queued behind step N's kernels that read it. The piecewise runner plans each region call the same way, just before its replay. A second slot only doubled every capture that touches the resource (the duplex encoder's 12 graphs became 24). Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/kv/bounded/manager.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/mstar/engine/resources/kv/bounded/manager.py b/mstar/engine/resources/kv/bounded/manager.py index e1abb47426..34e24bedfa 100644 --- a/mstar/engine/resources/kv/bounded/manager.py +++ b/mstar/engine/resources/kv/bounded/manager.py @@ -258,10 +258,11 @@ def commit(self, step: BoundedKVStep, ctx: StepContext) -> None: # from the declared position, so committing a step twice is harmless req.written = pos.written + seg.span - @property - def force_double_buffer(self): - return True - + # No ``force_double_buffer``: neither half of its hazard exists here. The host + # side is a PinnedStager ring, which never rewrites a buffer whose copy is in + # flight. And with no pre-plan, ``plan`` runs only on the stream that then + # runs the step, so plan(N+1)'s copy into a table queues behind step N's + # reads of it. A second slot would only double the captured graphs. def _table(self, ctx: StepContext, rows: int) -> torch.Tensor: lease = ctx.slot_lease if lease is None: From d67309c9f51c3c8d521b32a428f313199b8f0a5d Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:10:47 -0700 Subject: [PATCH 45/69] Bounded KV: tests for the per-query window, bf16, span 0 and the hooks The resource end to end (admit, plan, attend, commit) over a stream of steps against float64 brute force over each row's whole stream: per-query and per-step windows, at head_dim 128, in fp32 on the torch path and the kernel and in bf16 on the kernel; one-token steps and steps longer than 2 x BLOCK_N, rows padded past their span and rows that only read. Also: a read-only segment's commit after its region's leaves the count alone, an admit refusal names its request, the per-query config's checks, and dtype=None. The CUDA cases skip without a GPU. Not run here. Co-Authored-By: Claude Opus 5.5 --- test/modular/test_bounded_kv.py | 157 ++++++++++++++++++++++++++++++++ 1 file changed, 157 insertions(+) diff --git a/test/modular/test_bounded_kv.py b/test/modular/test_bounded_kv.py index d33c280973..cd73037840 100644 --- a/test/modular/test_bounded_kv.py +++ b/test/modular/test_bounded_kv.py @@ -256,3 +256,160 @@ def test_waypoint_ring_is_sink_window(ring_buckets, dilation): assert got == want, (frame, got, want) assert sum(c for _, c in layout.reads[:3]) == 0 written += span + + +# The resource end to end (layout, attention, writes) over a stream of steps, +# against brute force over each row's whole stream. + +CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernels need CUDA") +WINDOW = 70 + + +def _manager(window: int, per_query: bool, dtype, device, heads: int = 2, head_dim: int = 128): + from mstar.engine.resources.kv.bounded.config import BoundedKVConfig + from mstar.engine.resources.kv.bounded.manager import BoundedKVManager + + policy = SinkWindow(0, window, right=0 if per_query else None) + config = BoundedKVConfig(num_layers=1, num_heads=heads, head_dim=head_dim, rows_per_request=1, + max_source_len=0, retention=lambda n: policy, max_slots=4, dtype=dtype) + return BoundedKVManager(config, torch.device(device)) + + +def _step(rids, written, spans): + from mstar.engine.resources.kv.bounded.config import BoundedKVStep, StreamPosition + from mstar.engine.resources.step import Segment, StepContext + + step = BoundedKVStep( + segments=tuple(Segment(r, "main", s) for r, s in zip(rids, spans, strict=True)), + positions={r: StreamPosition(0, w) for r, w in zip(rids, written, strict=True)}, + ) + return step, StepContext(request_ids=tuple(rids), graph_walk="w", slot=0, capture=False) + + +def _stream_reference(streams, q, k, v, written, window, per_query, bias): + """float64: each row's query ``i`` (stream position ``w + i``) over its stream so + far ``streams[b] = (keys, values) [H, w, D]`` and the step's own tokens; per + query, the keys at distances [0, window], with the band indexed by distance; + per step, the last ``window`` stream keys and all of the step's own.""" + n, h, t, d = q.shape + out = torch.empty(n, t, h, d, dtype=torch.float64) + for b in range(n): + w = written[b] + keys = torch.cat([streams[b][0], k[b]], 1) + values = torch.cat([streams[b][1], v[b]], 1) + lo = max(0, w - window) + for i in range(t): + if per_query: + seen = torch.arange(max(0, w + i - window), w + i + 1) + rel = w + i - seen + else: + seen = torch.arange(lo, w + t) + rel = (w - lo) + i - (seen - lo) + t - 1 + scores = (q[b, :, i:i + 1] @ keys[:, seen].transpose(-1, -2))[:, 0] + bias[b, :, i, rel] + att = torch.softmax(scores * d ** -0.5, dim=-1) + out[b, i] = (att.unsqueeze(1) @ values[:, seen])[:, 0] + return out + + +@pytest.mark.parametrize("per_query", [True, False]) +@pytest.mark.parametrize("device,dtype", [ + ("cpu", torch.float32), + pytest.param("cuda", torch.float32, marks=CUDA), + pytest.param("cuda", torch.bfloat16, marks=CUDA), +]) +def test_stream_matches_brute_force(device, dtype, per_query): + """Rows at different positions in one batch, over one-token steps and steps + longer than 2 x BLOCK_N (where a per-query window leaves whole key tiles + masked for some queries), rows padded past their span, and rows that only + read (span 0), at head_dim 128.""" + torch.manual_seed(0) + h, d, window = 2, 128, WINDOW + mgr = _manager(window, per_query, dtype, device, heads=h, head_dim=d) + rids = [11, 12, 13] + streams = [(torch.zeros(h, 0, d, dtype=torch.float64),) * 2 for _ in rids] + written = [0, 0, 0] + # (T, each row's span) + schedule = [(1, [1, 1, 0]), (5, [5, 2, 0]), (150, [150, 0, 150]), (3, [3, 3, 1]), + (1, [1, 0, 1]), (200, [200, 137, 0]), (2, [2, 2, 2])] + for t, spans in schedule: + q, k, v = (torch.randn(len(rids), h, t, d).to(device=device, dtype=dtype) for _ in range(3)) + width = window + 1 if per_query else window + 2 * t - 1 + bias = (4 * torch.randn(len(rids), h, t, width)).to(device=device, dtype=dtype) + want = _stream_reference(streams, q.double().cpu(), k.double().cpu(), v.double().cpu(), + written, window, per_query, bias.double().cpu()) + step, ctx = _step(rids, written, spans) + assert mgr.admit(step, ctx).ok + mgr.plan(step, ctx) + got = mgr.attend(0, q, k, v, rel_bias=bias, ieee=True) + mgr.commit(step, ctx) + assert torch.isfinite(got.float()).all(), (t, spans) + tol = 2e-5 if dtype == torch.float32 else 3e-2 + torch.testing.assert_close(got.double().cpu(), want, rtol=tol, atol=tol) + for b, span in enumerate(spans): + sk, sv = streams[b] + streams[b] = (torch.cat([sk, k[b, :, :span].double().cpu()], 1), + torch.cat([sv, v[b, :, :span].double().cpu()], 1)) + written[b] += span + + +def test_commit_skips_read_only_segments(): + """An outer step that only reads a row commits after the region that wrote it, + and must not put the row's count back.""" + mgr = _manager(WINDOW, True, torch.float32, "cpu", heads=1, head_dim=4) + region, region_ctx = _step([1], [0], [1]) + outer, outer_ctx = _step([1], [0], [0]) + assert mgr.admit(outer, outer_ctx).ok + mgr.plan(outer, outer_ctx) + assert mgr.admit(region, region_ctx).ok + mgr.plan(region, region_ctx) + mgr.commit(region, region_ctx) + mgr.commit(outer, outer_ctx) + after, ctx = _step([1], [1], [1]) + assert mgr.admit(after, ctx).ok + + +def test_admit_refusal_names_the_request(): + mgr = _manager(WINDOW, True, torch.float32, "cpu", heads=1, head_dim=4) + first, ctx = _step([1, 2], [0, 0], [1, 1]) + assert mgr.admit(first, ctx).ok + mgr.plan(first, ctx) + mgr.commit(first, ctx) + # request 2 declares a stale position; request 1's is right + stale, ctx = _step([1, 2], [1, 0], [1, 1]) + outcome = mgr.admit(stale, ctx) + assert not outcome.ok and outcome.reason.request_ids == frozenset({2}) + + +def test_per_query_window_config(): + from mstar.engine.resources.kv.bounded.config import BoundedKVConfig + + def config(retention, **kw): + return BoundedKVConfig(num_layers=1, num_heads=1, head_dim=4, rows_per_request=1, + max_source_len=4, retention=retention, **kw) + + assert config(lambda n: SinkWindow(0, 70, right=0)).query_window == 70 + assert config(lambda n: SinkWindow(4, 70)).query_window is None + for bad in (SinkWindow(0, 70, right=1), SinkWindow(2, 70, right=0), + SinkWindow(0, 70, window_includes_step=True, right=0)): + with pytest.raises(ValueError): + config(lambda n, bad=bad: bad) + with pytest.raises(ValueError): + config(lambda n: SinkWindow(0, 70, right=0), reverse_step_order=True) + with pytest.raises(ValueError): # one kernel scalar for every source length + config(lambda n: SinkWindow(0, 70 + n, right=0)) + with pytest.raises(ValueError): + config(lambda n: SinkWindow(0, 70, right=0 if n else None)) + + +def test_dtype_none_is_the_engines_kv_dtype(): + from mstar.engine.resources.base import EngineResourceInfo, Resource + from mstar.engine.resources.kv.bounded.config import BoundedKVConfig, BoundedKVSpec + from mstar.engine.resources.kv.bounded.manager import BoundedKVManager + + config = BoundedKVConfig(num_layers=1, num_heads=1, head_dim=4, rows_per_request=1, max_source_len=0, + retention=lambda n: SinkWindow(0, 8, right=0), dtype=None) + spec = BoundedKVSpec(resource_key="kv", nodes={"node"}, config=config) + mgr = BoundedKVManager.build(spec, EngineResourceInfo(device=torch.device("cpu"), kv_dtype=torch.bfloat16)) + assert mgr.layer(0).dtype == torch.bfloat16 + # the engine sizes graph tables through main's hook name + assert BoundedKVManager.build_accelerator_graph_buffers is not Resource.build_accelerator_graph_buffers From 94e5067b49424bbdc1d134bb36291868f418cb1f Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:11:03 -0700 Subject: [PATCH 46/69] docs: bounded KV's per-query window and dtype=None Co-Authored-By: Claude Opus 5.5 --- docs/adding_models.rst | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/docs/adding_models.rst b/docs/adding_models.rst index b4ce9774b8..b73488aec8 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -711,6 +711,16 @@ Two flags cover retention rules that are not quite a token FIFO. only makes sense for attention that ignores key order; it decides which of the oldest step's tokens leave first. +By default the window is per step: every query of a step sees the same keys, the retained +ones and all of the step's own. A streaming encoder whose frames each see a fixed left +context (a conformer's ``att_context_size=[left, 0]``) sets ``SinkWindow(0, left, right=0)`` +instead: each query attends the keys at distances ``0..left`` from it, so a step of several +frames is causal within itself and matches the offline mask. ``rel_bias`` is then indexed +by that distance, ``[rows, H, T, left + 1]``, and passed unscaled (the kernel scales the +sum by ``head_dim ** -0.5``). Only a right context of 0, with no sink, is supported so far. +``BoundedKVConfig(dtype=None)`` stores the cache in the engine's KV dtype, for layers that +compute in it. + Step 3 — Declare the computation graph -------------------------------------- From de6be73054c00a016a29ec66ce5bfb3ce3bc0821 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:52:09 -0700 Subject: [PATCH 47/69] Bounded KV: the per-query window is a config field, retention may be constant A per-query window is a property of the cache's attention, not of what a source length retains, so it moves from SinkWindow.right to BoundedKVConfig.query_window (None keeps the per-step window). Its restrictions stay config-time checks; the check that every source length picks the same window goes, since there is now one. retention may also be a single SinkWindow, so two configs with the same constant policy compare equal by value; the callable form is unchanged. Co-Authored-By: Claude Opus 5.5 --- docs/adding_models.rst | 12 +++-- mstar/engine/resources/kv/bounded/config.py | 49 ++++++++------------ mstar/engine/resources/kv/bounded/kernels.py | 3 +- mstar/engine/resources/kv/bounded/manager.py | 4 +- test/modular/test_bounded_kv.py | 38 +++++++++------ 5 files changed, 55 insertions(+), 51 deletions(-) diff --git a/docs/adding_models.rst b/docs/adding_models.rst index b73488aec8..5f0a23eeb1 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -713,11 +713,13 @@ step's tokens leave first. By default the window is per step: every query of a step sees the same keys, the retained ones and all of the step's own. A streaming encoder whose frames each see a fixed left -context (a conformer's ``att_context_size=[left, 0]``) sets ``SinkWindow(0, left, right=0)`` -instead: each query attends the keys at distances ``0..left`` from it, so a step of several -frames is causal within itself and matches the offline mask. ``rel_bias`` is then indexed -by that distance, ``[rows, H, T, left + 1]``, and passed unscaled (the kernel scales the -sum by ``head_dim ** -0.5``). Only a right context of 0, with no sink, is supported so far. +context (a conformer's ``att_context_size=[left, 0]``) sets +``BoundedKVConfig(retention=SinkWindow(0, left), query_window=left)`` instead: each query +attends the keys at distances ``0..left`` from it, so a step of several frames is causal +within itself and matches the offline mask. ``rel_bias`` is then indexed by that distance, +``[rows, H, T, left + 1]``, and passed unscaled (the kernel scales the sum by +``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. diff --git a/mstar/engine/resources/kv/bounded/config.py b/mstar/engine/resources/kv/bounded/config.py index 8b83000976..e8c9ededeb 100644 --- a/mstar/engine/resources/kv/bounded/config.py +++ b/mstar/engine/resources/kv/bounded/config.py @@ -25,19 +25,11 @@ class SinkWindow(NamedTuple): """Keep the stream's first ``sink`` and last ``window`` entries. With ``window_includes_step`` the window counts the step's own tokens, so a step - that writes ``span`` attends only the last ``window - span`` retained ones. - - By default the window is per step: every query of a step attends the same - keys, the retained ones and all of the step's own. With ``right`` set it is - per query: each query attends the keys at distances ``[-right, window]`` - from it, so a step of several tokens is causal within itself and each of - them sees exactly ``window`` back (a streaming conformer's - ``att_context_size``). Only ``right=0``, with no sink, is supported so far.""" + that writes ``span`` attends only the last ``window - span`` retained ones.""" sink: int window: int window_includes_step: bool = False - right: int | None = None @dataclass @@ -49,8 +41,8 @@ class BoundedKVConfig: rows_per_request: int # the longest source prefix a request may have; sizes the slots max_source_len: int - # source length -> what that request keeps - retention: Callable[[int], SinkWindow] + # source length -> what that request keeps, or one policy for every length + retention: Callable[[int], SinkWindow] | SinkWindow max_slots: int = 32 # the shortest source prefix a request may have: the slots also hold the part # of a sink its source doesn't cover, which is largest for the shortest source @@ -62,34 +54,38 @@ class BoundedKVConfig: # ignores key order, where it decides which of the oldest step's tokens are # evicted first: Step-Audio2's flow decoder keeps the front of each chunk. reverse_step_order: bool = False + # None: the window is per step, every query of a step attending the same keys + # (the retained ones and all of the step's own). Set, it is per query: each + # query attends the keys at distances [0, query_window] from it, so a step of + # several tokens is causal within itself (a streaming conformer's + # ``att_context_size=[left, 0]``). Needs no sink and the stream in order. + query_window: int | None = None def __post_init__(self): if self.max_slots < 1: raise ValueError(f"max_slots must be >= 1, got {self.max_slots}") if not 0 <= self.min_source_len <= self.max_source_len: raise ValueError(f"source lengths {self.min_source_len}..{self.max_source_len} are not a range") + q_window = self.query_window + if q_window is not None and (q_window < 0 or self.reverse_step_order): + raise ValueError(f"query_window={q_window} needs to be >= 0, with the stream in order") sink = window = retained = 0 - query_windows = set() for n in range(self.min_source_len, self.max_source_len + 1): - policy = self.retention(n) + policy = self.retention_for(n) if policy.sink < 0 or policy.window < 0: raise ValueError(f"retention for {n} source tokens is {policy}") - if policy.right is not None and ( - policy.right != 0 or policy.sink or policy.window_includes_step or self.reverse_step_order - ): - raise ValueError(f"retention for {n} source tokens is {policy}: a per-query window " - "takes right=0, no sink, and the stream in order") - query_windows.add(None if policy.right is None else policy.window) + if q_window is not None and (policy.sink or policy.window_includes_step or policy.window < q_window): + raise ValueError(f"retention for {n} source tokens is {policy}: a per-query window of " + f"{q_window} takes no sink and a window of at least {q_window}") sink = max(sink, policy.sink - n) window = max(window, policy.window) retained = max(retained, policy.sink + policy.window) - # the kernel takes a per-query window as one scalar for the whole step - if len(query_windows) > 1: - raise ValueError(f"retention mixes windows {sorted(query_windows, key=str)}; a per-query " - "window must be the same for every source length") # fixed once here: a plan reads them for every row self._sink_capacity, self._window_capacity, self._max_retained = sink, window, retained - self._query_window = next(iter(query_windows)) + + def retention_for(self, source_len: int) -> SinkWindow: + r = self.retention + return r if isinstance(r, SinkWindow) else r(source_len) @property def sink_capacity(self) -> int: @@ -100,11 +96,6 @@ def sink_capacity(self) -> int: def window_capacity(self) -> int: return self._window_capacity - @property - def query_window(self) -> int | None: - """Each query's window when retention is per query (``SinkWindow.right``), else None.""" - return self._query_window - @property def max_retained(self) -> int: """The most entries a request retains, over every allowed source length.""" diff --git a/mstar/engine/resources/kv/bounded/kernels.py b/mstar/engine/resources/kv/bounded/kernels.py index 02b3c874b2..e000fc6c68 100644 --- a/mstar/engine/resources/kv/bounded/kernels.py +++ b/mstar/engine/resources/kv/bounded/kernels.py @@ -3,7 +3,8 @@ Keys and values share a ``2 * head_dim`` last axis in the slot and the source: keys in ``[0, D)``, values in ``[D, 2D)``. Every row attends to its retained ranges and to this step's own tokens: unmasked, or with a per-query window -(``SinkWindow.right``) each query to the keys at distances ``[0, window]``. +(``BoundedKVConfig.query_window``) each query to the keys at distances +``[0, window]``. Dot products use TF32 unless ``ieee``. """ diff --git a/mstar/engine/resources/kv/bounded/manager.py b/mstar/engine/resources/kv/bounded/manager.py index 34e24bedfa..44f2fb3d2a 100644 --- a/mstar/engine/resources/kv/bounded/manager.py +++ b/mstar/engine/resources/kv/bounded/manager.py @@ -107,7 +107,7 @@ def attend( step's own, then this step's keys into the slots. ``[B * rows, T, H, D]``. ``plan`` defaults to the step's. ``rel_bias`` is a per-distance score term (``kernels.bounded_attention``), for relative-position attention; with a - per-query window (``SinkWindow.right``) it is indexed by distance, + per-query window (``BoundedKVConfig.query_window``) it is indexed by distance, ``[B * rows, H, T, window + 1]``. Rows whose source is -1 read theirs from ``row_source = (storage [slots, rows, H, L, 2D], slots [>= B])``, e.g. another resource's per-request entries.""" @@ -235,7 +235,7 @@ def plan(self, step: BoundedKVStep, ctx: StepContext) -> BoundedPlan: else: pos = step.positions[seg.request_id] rows.append((req.slot, row_layout(pos.source_len, pos.written, seg.span, - c.retention(pos.source_len), c.sink_capacity, + c.retention_for(pos.source_len), c.sink_capacity, c.reverse_step_order), pos.source)) values += flatten_row(*rows[-1]) table = self._table(ctx, len(rows)) diff --git a/test/modular/test_bounded_kv.py b/test/modular/test_bounded_kv.py index cd73037840..25c507e5ed 100644 --- a/test/modular/test_bounded_kv.py +++ b/test/modular/test_bounded_kv.py @@ -76,7 +76,7 @@ def test_slots_hold_the_sink_of_the_shortest_source(source_len): ) assert config.sink_capacity == 30 and config.max_retained == 50 source = [("s", i) for i in range(source_len)] - sim = SlotSim(source, config.retention(source_len), config.sink_capacity) + sim = SlotSim(source, config.retention_for(source_len), config.sink_capacity) stream = list(source) for step in range(8): fresh = [("g", len(stream) - source_len + j) for j in range(7)] @@ -269,9 +269,9 @@ def _manager(window: int, per_query: bool, dtype, device, heads: int = 2, head_d from mstar.engine.resources.kv.bounded.config import BoundedKVConfig from mstar.engine.resources.kv.bounded.manager import BoundedKVManager - policy = SinkWindow(0, window, right=0 if per_query else None) config = BoundedKVConfig(num_layers=1, num_heads=heads, head_dim=head_dim, rows_per_request=1, - max_source_len=0, retention=lambda n: policy, max_slots=4, dtype=dtype) + max_source_len=0, retention=SinkWindow(0, window), max_slots=4, dtype=dtype, + query_window=window if per_query else None) return BoundedKVManager(config, torch.device(device)) @@ -387,18 +387,28 @@ def config(retention, **kw): return BoundedKVConfig(num_layers=1, num_heads=1, head_dim=4, rows_per_request=1, max_source_len=4, retention=retention, **kw) - assert config(lambda n: SinkWindow(0, 70, right=0)).query_window == 70 - assert config(lambda n: SinkWindow(4, 70)).query_window is None - for bad in (SinkWindow(0, 70, right=1), SinkWindow(2, 70, right=0), - SinkWindow(0, 70, window_includes_step=True, right=0)): + assert config(lambda n: SinkWindow(0, 70), query_window=70).query_window == 70 + # a longer retained window than the query's is allowed (its far keys are masked) + config(lambda n: SinkWindow(0, 70 + n), query_window=70) + for bad in (SinkWindow(2, 70), SinkWindow(0, 70, window_includes_step=True), SinkWindow(0, 69)): with pytest.raises(ValueError): - config(lambda n, bad=bad: bad) + config(bad, query_window=70) with pytest.raises(ValueError): - config(lambda n: SinkWindow(0, 70, right=0), reverse_step_order=True) - with pytest.raises(ValueError): # one kernel scalar for every source length - config(lambda n: SinkWindow(0, 70 + n, right=0)) - with pytest.raises(ValueError): - config(lambda n: SinkWindow(0, 70, right=0 if n else None)) + config(SinkWindow(0, 70), query_window=70, reverse_step_order=True) + + +def test_constant_retention_compares_by_value(): + """A constant policy works like the callable form, and two configs built from + equal ones are equal (a model's spec compared against another's).""" + from mstar.engine.resources.kv.bounded.config import BoundedKVConfig + + def config(retention): + return BoundedKVConfig(num_layers=1, num_heads=1, head_dim=4, rows_per_request=1, + max_source_len=4, retention=retention, query_window=70) + + assert config(SinkWindow(0, 70)) == config(SinkWindow(0, 70)) + assert config(SinkWindow(0, 70)) != config(SinkWindow(0, 71)) + assert config(SinkWindow(0, 70)).retention_for(3) == config(lambda n: SinkWindow(0, 70)).retention_for(3) def test_dtype_none_is_the_engines_kv_dtype(): @@ -407,7 +417,7 @@ def test_dtype_none_is_the_engines_kv_dtype(): from mstar.engine.resources.kv.bounded.manager import BoundedKVManager config = BoundedKVConfig(num_layers=1, num_heads=1, head_dim=4, rows_per_request=1, max_source_len=0, - retention=lambda n: SinkWindow(0, 8, right=0), dtype=None) + retention=SinkWindow(0, 8), dtype=None) spec = BoundedKVSpec(resource_key="kv", nodes={"node"}, config=config) mgr = BoundedKVManager.build(spec, EngineResourceInfo(device=torch.device("cpu"), kv_dtype=torch.bfloat16)) assert mgr.layer(0).dtype == torch.bfloat16 From 6cc08d0d3368449cecb3fa3f011d9553f5a7b6a9 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:52:22 -0700 Subject: [PATCH 48/69] Bounded KV: resolve retention once per source length plan called the retention callable for every row of every step; the config already evaluates it for each allowed source length when it is built, so it keeps those and plan looks one up. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/kv/bounded/config.py | 11 +++++----- test/modular/test_bounded_kv.py | 23 +++++++++++++++++++++ 2 files changed, 29 insertions(+), 5 deletions(-) diff --git a/mstar/engine/resources/kv/bounded/config.py b/mstar/engine/resources/kv/bounded/config.py index e8c9ededeb..385f54d870 100644 --- a/mstar/engine/resources/kv/bounded/config.py +++ b/mstar/engine/resources/kv/bounded/config.py @@ -69,9 +69,12 @@ def __post_init__(self): q_window = self.query_window if q_window is not None and (q_window < 0 or self.reverse_step_order): raise ValueError(f"query_window={q_window} needs to be >= 0, with the stream in order") + r = self.retention + # resolved once here: a plan looks one up for every row + self._policies = tuple(r if isinstance(r, SinkWindow) else r(n) + for n in range(self.min_source_len, self.max_source_len + 1)) sink = window = retained = 0 - for n in range(self.min_source_len, self.max_source_len + 1): - policy = self.retention_for(n) + for n, policy in enumerate(self._policies, self.min_source_len): if policy.sink < 0 or policy.window < 0: raise ValueError(f"retention for {n} source tokens is {policy}") if q_window is not None and (policy.sink or policy.window_includes_step or policy.window < q_window): @@ -80,12 +83,10 @@ def __post_init__(self): sink = max(sink, policy.sink - n) window = max(window, policy.window) retained = max(retained, policy.sink + policy.window) - # fixed once here: a plan reads them for every row self._sink_capacity, self._window_capacity, self._max_retained = sink, window, retained def retention_for(self, source_len: int) -> SinkWindow: - r = self.retention - return r if isinstance(r, SinkWindow) else r(source_len) + return self._policies[source_len - self.min_source_len] @property def sink_capacity(self) -> int: diff --git a/test/modular/test_bounded_kv.py b/test/modular/test_bounded_kv.py index 25c507e5ed..849aec2f34 100644 --- a/test/modular/test_bounded_kv.py +++ b/test/modular/test_bounded_kv.py @@ -411,6 +411,29 @@ def config(retention): assert config(SinkWindow(0, 70)).retention_for(3) == config(lambda n: SinkWindow(0, 70)).retention_for(3) +def test_plan_resolves_retention_once_per_source_len(): + """The policy is resolved when the config is built, not per row per step.""" + from mstar.engine.resources.kv.bounded.config import BoundedKVConfig + from mstar.engine.resources.kv.bounded.manager import BoundedKVManager + + calls = [] + + def retention(n): + calls.append(n) + return SinkWindow(0, WINDOW) + + config = BoundedKVConfig(num_layers=1, num_heads=1, head_dim=4, rows_per_request=1, max_source_len=0, + retention=retention, max_slots=4) + mgr = BoundedKVManager(config, torch.device("cpu")) + built = len(calls) + for t in range(3): + step, ctx = _step([1, 2], [t, t], [1, 1]) + assert mgr.admit(step, ctx).ok + mgr.plan(step, ctx) + mgr.commit(step, ctx) + assert len(calls) == built + + def test_dtype_none_is_the_engines_kv_dtype(): from mstar.engine.resources.base import EngineResourceInfo, Resource from mstar.engine.resources.kv.bounded.config import BoundedKVConfig, BoundedKVSpec From e03a4a286be11209887c06c4e84a37288bcc180f Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:53:12 -0700 Subject: [PATCH 49/69] Bounded KV: per step, a row attends only the step's keys it writes The own-key loop ran to T, so a row padded past its span (a shorter row in a longer step) attended its padding keys, and the CPU path mirrored it. The row table now carries each row's span, and per step a row attends the first span of the step's keys, or all T when it writes none (a read-only row's tokens are all real). A per-query window needs nothing: causality already keeps a row's real queries off its padding. The stream test's reference had the bug baked in; it now attends the span. Co-Authored-By: Claude Opus 5.5 --- docs/adding_models.rst | 5 +++-- mstar/engine/resources/kv/bounded/kernels.py | 20 +++++++++++--------- mstar/engine/resources/kv/bounded/layout.py | 12 ++++++++---- mstar/engine/resources/kv/bounded/manager.py | 5 +++-- test/modular/test_bounded_kv.py | 10 ++++++---- 5 files changed, 31 insertions(+), 21 deletions(-) diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 5f0a23eeb1..4fe7b62ebd 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -697,7 +697,8 @@ computed on the host from the position the step declares, which is also what let different positions share a batch and a captured graph. The model calls ``attend(layer_idx, q, k, v, source, plan)``: each row attends its -retained keys in stream order and its own, unmasked, and its writes land after. Pass +retained keys in stream order and the step's keys it writes (all of them for a row that +writes none), unmasked, and its writes land after. Pass ``source`` (``[sources, rows, H, L, 2 * head_dim]``, keys then values) when requests start from shared prefixes that the slots should not copy; each row reads the one its ``StreamPosition.source`` names, so rows of different prefixes share a batch and a graph. ``rel_bias`` adds a term that depends only on @@ -712,7 +713,7 @@ only makes sense for attention that ignores key order; it decides which of the o step's tokens leave first. By default the window is per step: every query of a step sees the same keys, the retained -ones and all of the step's own. A streaming encoder whose frames each see a fixed left +ones and the step's own. A streaming encoder whose frames each see a fixed left context (a conformer's ``att_context_size=[left, 0]``) sets ``BoundedKVConfig(retention=SinkWindow(0, left), query_window=left)`` instead: each query attends the keys at distances ``0..left`` from it, so a step of several frames is causal diff --git a/mstar/engine/resources/kv/bounded/kernels.py b/mstar/engine/resources/kv/bounded/kernels.py index e000fc6c68..c6d378fc92 100644 --- a/mstar/engine/resources/kv/bounded/kernels.py +++ b/mstar/engine/resources/kv/bounded/kernels.py @@ -2,10 +2,10 @@ Keys and values share a ``2 * head_dim`` last axis in the slot and the source: keys in ``[0, D)``, values in ``[D, 2D)``. Every row attends to its retained -ranges and to this step's own tokens: unmasked, or with a per-query window -(``BoundedKVConfig.query_window``) each query to the keys at distances -``[0, window]``. -Dot products use TF32 unless ``ieee``. +ranges and to this step's own tokens: unmasked (those it writes, or all T when +it writes none), or with a per-query window (``BoundedKVConfig.query_window``) +each query to the keys at distances ``[0, window]``. Dot products use TF32 +unless ``ieee``. """ from __future__ import annotations @@ -92,16 +92,18 @@ def _bounded_attention_kernel( own_hi = tl.minimum(m_lo + BLOCK_M, T) old_lo = tl.maximum(total + m_lo - window, 0) else: + own = tl.load(row + 2) own_lo = 0 - own_hi = T + own_hi = tl.where(own > 0, tl.minimum(own, T), T) old_lo = 0 - # this step's own tokens (separate k and v tensors, laid out like q); a pair's - # band entry is its query position minus its key position, plus T - 1, or - # with a per-query window the distance alone + # this step's own tokens (separate k and v tensors, laid out like q), per step + # the first ``span`` (``RowLayout.span``); a pair's band entry is its query + # position minus its key position, plus T - 1, or with a per-query window the + # distance alone (causality already keeps a row's real queries off its padding) for j0 in range(own_lo, own_hi, BLOCK_N): offs_n = j0 + tl.arange(0, BLOCK_N) - mask_n = offs_n < T + mask_n = offs_n < own_hi kv_off = n * s_qn + h * s_qh + offs_n[:, None] * s_qt + offs_d[None, :] k = tl.load(k_ptr + kv_off, mask=mask_n[:, None], other=0.0) v = tl.load(v_ptr + kv_off, mask=mask_n[:, None], other=0.0) diff --git a/mstar/engine/resources/kv/bounded/layout.py b/mstar/engine/resources/kv/bounded/layout.py index 4c511f46a1..035573b971 100644 --- a/mstar/engine/resources/kv/bounded/layout.py +++ b/mstar/engine/resources/kv/bounded/layout.py @@ -18,14 +18,18 @@ READ_FROM_SOURCE = (True, False, True, False, False) # writes: (fresh offset of the range's first entry, slot start, len) triples NUM_WRITES = 3 -# slot, source, reads, writes -ROW_HEAD = 2 +# slot, source, span, reads, writes +ROW_HEAD = 3 ROW_INTS = ROW_HEAD + 2 * NUM_READS + 3 * NUM_WRITES class RowLayout(NamedTuple): reads: tuple[tuple[int, int], ...] writes: tuple[tuple[int, int, int], ...] + # what the row writes this step; per step it also attends only those of the + # step's keys (a row padded past its span must not see its padding), or all + # of them when it writes none + span: int = 0 def _ring_ranges(first: int, count: int, window: int) -> list[tuple[int, int]]: @@ -71,14 +75,14 @@ def fresh(i: int) -> int: writes.append((fresh(first), ring + s, c)) first += c writes += [(0, 0, 0)] * (NUM_WRITES - len(writes)) - return RowLayout(tuple(reads), tuple(writes)) + return RowLayout(tuple(reads), tuple(writes), span) EMPTY_ROW = RowLayout(((0, 0),) * NUM_READS, ((0, 0, 0),) * NUM_WRITES) def flatten_row(slot: int, layout: RowLayout, source: int = 0) -> list[int]: - out = [slot, source] + out = [slot, source, layout.span] for start, count in layout.reads: out += (start, count) for offset, start, count in layout.writes: diff --git a/mstar/engine/resources/kv/bounded/manager.py b/mstar/engine/resources/kv/bounded/manager.py index 44f2fb3d2a..1442ba78fc 100644 --- a/mstar/engine/resources/kv/bounded/manager.py +++ b/mstar/engine/resources/kv/bounded/manager.py @@ -143,10 +143,11 @@ def _attend_torch( kv = (row_src if from_source else cache[slot, c])[:, start:start + count] keys.append(kv[..., :d]) values.append(kv[..., d:]) - keys, values = torch.cat([*keys, k[i]], 1), torch.cat([*values, v[i]], 1) + total = sum(part.shape[1] for part in keys) + own_len = t if window is not None or layout.span == 0 else min(layout.span, t) + keys, values = torch.cat([*keys, k[i][:, :own_len]], 1), torch.cat([*values, v[i][:, :own_len]], 1) scores = q[i] @ keys.transpose(-1, -2) # query i sits at total + i, key j at j - total = keys.shape[1] - t dist = (total + torch.arange(t)[:, None] - torch.arange(keys.shape[1])[None, :]).to(q.device) if window is None: if rel_bias is not None: diff --git a/test/modular/test_bounded_kv.py b/test/modular/test_bounded_kv.py index 849aec2f34..96ee9e794f 100644 --- a/test/modular/test_bounded_kv.py +++ b/test/modular/test_bounded_kv.py @@ -286,11 +286,12 @@ def _step(rids, written, spans): return step, StepContext(request_ids=tuple(rids), graph_walk="w", slot=0, capture=False) -def _stream_reference(streams, q, k, v, written, window, per_query, bias): +def _stream_reference(streams, q, k, v, written, spans, window, per_query, bias): """float64: each row's query ``i`` (stream position ``w + i``) over its stream so far ``streams[b] = (keys, values) [H, w, D]`` and the step's own tokens; per query, the keys at distances [0, window], with the band indexed by distance; - per step, the last ``window`` stream keys and all of the step's own.""" + per step, the last ``window`` stream keys and the step's first ``span`` (all T + for a row that writes none).""" n, h, t, d = q.shape out = torch.empty(n, t, h, d, dtype=torch.float64) for b in range(n): @@ -298,12 +299,13 @@ def _stream_reference(streams, q, k, v, written, window, per_query, bias): keys = torch.cat([streams[b][0], k[b]], 1) values = torch.cat([streams[b][1], v[b]], 1) lo = max(0, w - window) + own = spans[b] or t for i in range(t): if per_query: seen = torch.arange(max(0, w + i - window), w + i + 1) rel = w + i - seen else: - seen = torch.arange(lo, w + t) + seen = torch.arange(lo, w + own) rel = (w - lo) + i - (seen - lo) + t - 1 scores = (q[b, :, i:i + 1] @ keys[:, seen].transpose(-1, -2))[:, 0] + bias[b, :, i, rel] att = torch.softmax(scores * d ** -0.5, dim=-1) @@ -336,7 +338,7 @@ def test_stream_matches_brute_force(device, dtype, per_query): width = window + 1 if per_query else window + 2 * t - 1 bias = (4 * torch.randn(len(rids), h, t, width)).to(device=device, dtype=dtype) want = _stream_reference(streams, q.double().cpu(), k.double().cpu(), v.double().cpu(), - written, window, per_query, bias.double().cpu()) + written, spans, window, per_query, bias.double().cpu()) step, ctx = _step(rids, written, spans) assert mgr.admit(step, ctx).ok mgr.plan(step, ctx) From 2acf397f486ffed3c9af79b9bb45c609657e2562 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:54:51 -0700 Subject: [PATCH 50/69] Bounded KV: pick the query tile from T and the element size BLOCK_M was 64 whenever n * h >= 128, whatever the dtype and step length. At head_dim 128 in fp32 that tile needs 245760 (per-query) or 294912 (per-step) bytes of shared memory under TF32, past H100's 232448, and under IEEE it compiles but spills 2686-4592 registers; at 16 it takes ~143 KB with no spills. bf16 at D=128 fits 64 (120-147 KB, 16 spills at most). So 64 needs a step longer than 16 queries and at most 256 bytes per head row. The new test runs 32 rows x 8 heads at D=128 in fp32 and bf16, ieee on and off, T 1 and 150, both window modes, against float64 brute force. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/kv/bounded/kernels.py | 6 +- test/modular/test_bounded_kv.py | 72 +++++++++++++------- 2 files changed, 52 insertions(+), 26 deletions(-) diff --git a/mstar/engine/resources/kv/bounded/kernels.py b/mstar/engine/resources/kv/bounded/kernels.py index c6d378fc92..4558339bda 100644 --- a/mstar/engine/resources/kv/bounded/kernels.py +++ b/mstar/engine/resources/kv/bounded/kernels.py @@ -242,7 +242,11 @@ def bounded_attention( else: rel_bias, s_bias = q, (0, 0, 0) out = torch.empty(n, t, h, d, device=q.device, dtype=q.dtype) - block_m = 16 if n * h < 128 else 64 + # 64-query tiles only where they pay and fit: a short step wastes them, too few + # programs idle SMs, and past 256 bytes per head row (fp32 at D=128) one + # overflows shared memory under TF32 and spills thousands of registers under + # IEEE (H100), where a 16-query tile fits with no spills. + block_m = 64 if t > 16 and n * h >= 128 and d * q.element_size() <= 256 else 16 grid = (triton.cdiv(t, block_m), n * h) _bounded_attention_kernel[grid]( q, k, v, out, cache, source, row_storage, row_slots, table, rel_bias, diff --git a/test/modular/test_bounded_kv.py b/test/modular/test_bounded_kv.py index 96ee9e794f..5aa69549a7 100644 --- a/test/modular/test_bounded_kv.py +++ b/test/modular/test_bounded_kv.py @@ -265,12 +265,12 @@ def test_waypoint_ring_is_sink_window(ring_buckets, dilation): WINDOW = 70 -def _manager(window: int, per_query: bool, dtype, device, heads: int = 2, head_dim: int = 128): +def _manager(window: int, per_query: bool, dtype, device, heads: int = 2, head_dim: int = 128, slots: int = 4): from mstar.engine.resources.kv.bounded.config import BoundedKVConfig from mstar.engine.resources.kv.bounded.manager import BoundedKVManager config = BoundedKVConfig(num_layers=1, num_heads=heads, head_dim=head_dim, rows_per_request=1, - max_source_len=0, retention=SinkWindow(0, window), max_slots=4, dtype=dtype, + max_source_len=0, retention=SinkWindow(0, window), max_slots=slots, dtype=dtype, query_window=window if per_query else None) return BoundedKVManager(config, torch.device(device)) @@ -313,39 +313,27 @@ def _stream_reference(streams, q, k, v, written, spans, window, per_query, bias) return out -@pytest.mark.parametrize("per_query", [True, False]) -@pytest.mark.parametrize("device,dtype", [ - ("cpu", torch.float32), - pytest.param("cuda", torch.float32, marks=CUDA), - pytest.param("cuda", torch.bfloat16, marks=CUDA), -]) -def test_stream_matches_brute_force(device, dtype, per_query): - """Rows at different positions in one batch, over one-token steps and steps - longer than 2 x BLOCK_N (where a per-query window leaves whole key tiles - masked for some queries), rows padded past their span, and rows that only - read (span 0), at head_dim 128.""" - torch.manual_seed(0) - h, d, window = 2, 128, WINDOW - mgr = _manager(window, per_query, dtype, device, heads=h, head_dim=d) - rids = [11, 12, 13] +def _check_stream(mgr, schedule, ieee, tol): + """Runs ``schedule`` [(T, each row's span)] through ``mgr`` against brute force.""" + c = mgr.config + h, d, window, per_query = c.num_heads, c.head_dim, c.window_capacity, c.query_window is not None + device, dtype = mgr.layer(0).device, mgr.layer(0).dtype + rows = len(schedule[0][1]) + rids = list(range(100, 100 + rows)) streams = [(torch.zeros(h, 0, d, dtype=torch.float64),) * 2 for _ in rids] - written = [0, 0, 0] - # (T, each row's span) - schedule = [(1, [1, 1, 0]), (5, [5, 2, 0]), (150, [150, 0, 150]), (3, [3, 3, 1]), - (1, [1, 0, 1]), (200, [200, 137, 0]), (2, [2, 2, 2])] + written = [0] * rows for t, spans in schedule: - q, k, v = (torch.randn(len(rids), h, t, d).to(device=device, dtype=dtype) for _ in range(3)) + q, k, v = (torch.randn(rows, h, t, d).to(device=device, dtype=dtype) for _ in range(3)) width = window + 1 if per_query else window + 2 * t - 1 - bias = (4 * torch.randn(len(rids), h, t, width)).to(device=device, dtype=dtype) + bias = (4 * torch.randn(rows, h, t, width)).to(device=device, dtype=dtype) want = _stream_reference(streams, q.double().cpu(), k.double().cpu(), v.double().cpu(), written, spans, window, per_query, bias.double().cpu()) step, ctx = _step(rids, written, spans) assert mgr.admit(step, ctx).ok mgr.plan(step, ctx) - got = mgr.attend(0, q, k, v, rel_bias=bias, ieee=True) + got = mgr.attend(0, q, k, v, rel_bias=bias, ieee=ieee) mgr.commit(step, ctx) assert torch.isfinite(got.float()).all(), (t, spans) - tol = 2e-5 if dtype == torch.float32 else 3e-2 torch.testing.assert_close(got.double().cpu(), want, rtol=tol, atol=tol) for b, span in enumerate(spans): sk, sv = streams[b] @@ -354,6 +342,40 @@ def test_stream_matches_brute_force(device, dtype, per_query): written[b] += span +@pytest.mark.parametrize("per_query", [True, False]) +@pytest.mark.parametrize("device,dtype", [ + ("cpu", torch.float32), + pytest.param("cuda", torch.float32, marks=CUDA), + pytest.param("cuda", torch.bfloat16, marks=CUDA), +]) +def test_stream_matches_brute_force(device, dtype, per_query): + """Rows at different positions in one batch, over one-token steps and steps + longer than 2 x BLOCK_N (where a per-query window leaves whole key tiles + masked for some queries), rows padded past their span, and rows that only + read (span 0), at head_dim 128.""" + torch.manual_seed(0) + mgr = _manager(WINDOW, per_query, dtype, device, heads=2, head_dim=128) + schedule = [(1, [1, 1, 0]), (5, [5, 2, 0]), (150, [150, 0, 150]), (3, [3, 3, 1]), + (1, [1, 0, 1]), (200, [200, 137, 0]), (2, [2, 2, 2])] + _check_stream(mgr, schedule, ieee=True, tol=2e-5 if dtype == torch.float32 else 3e-2) + + +@CUDA +@pytest.mark.parametrize("t", [1, 150]) +@pytest.mark.parametrize("per_query", [True, False]) +@pytest.mark.parametrize("dtype,ieee,tol", [ + (torch.float32, True, 2e-5), (torch.float32, False, 2e-3), + (torch.bfloat16, True, 3e-2), (torch.bfloat16, False, 3e-2), +]) +def test_stream_at_many_rows(dtype, ieee, tol, per_query, t): + """32 rows x 8 heads (enough programs for the kernel's wide query tiles) at + head_dim 128: every tile config a real model reaches compiles and is right.""" + torch.manual_seed(0) + mgr = _manager(WINDOW, per_query, dtype, "cuda", heads=8, head_dim=128, slots=32) + spans = [1, 1, 0, 1] if t == 1 else [150, 0, 99, 150] + _check_stream(mgr, [(150, [150, 37, 0, 150] * 8), (t, spans * 8)], ieee=ieee, tol=tol) + + def test_commit_skips_read_only_segments(): """An outer step that only reads a row commits after the region that wrote it, and must not put the row's count back.""" From c0e4499ad86ce21c6a5884e89992af88f054af02 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:57:34 -0700 Subject: [PATCH 51/69] Bounded KV: rollback_admit gives back the slots an admit took admit takes a slot for a request's first step, and a refused step (here or in another resource) kept it, contrary to the base contract. Like RecurrentStatePool, admit now records the slots it takes fresh and rollback_admit frees them; a slot a request already held stays, and capture and padding rows take nothing to give back. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/kv/bounded/manager.py | 22 ++++++++++++++++++++ test/modular/test_bounded_kv.py | 22 ++++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/mstar/engine/resources/kv/bounded/manager.py b/mstar/engine/resources/kv/bounded/manager.py index 1442ba78fc..d446f11748 100644 --- a/mstar/engine/resources/kv/bounded/manager.py +++ b/mstar/engine/resources/kv/bounded/manager.py @@ -70,6 +70,8 @@ def __init__(self, config: BoundedKVConfig, device: torch.device): self._requests: dict[str, _Request] = {} self._free = list(reversed(range(c.max_slots))) self._lock = threading.RLock() + # rid -> slot the live admit took fresh, for ``rollback_admit`` + self._admit_taken: dict[str, int] = {} self._stager = PinnedStager(torch.int32) self._cg_max_bs = 0 @@ -188,6 +190,7 @@ def reset_request(self, rid: str, free: bool = False): return if free: del self._requests[rid] + self._admit_taken.pop(rid, None) self._free.append(req.slot) else: req.written = 0 @@ -198,6 +201,9 @@ def admit(self, step: BoundedKVStep, ctx: StepContext) -> AdmitOutcome: if ctx.capture: return ADMIT_OK with self._lock: + taken = self._admit_taken + for rid in [r for r in taken if ctx.is_real_row(r)]: + del taken[rid] for seg in step.segments or (): if ctx.is_padding_row(seg.request_id): continue @@ -218,12 +224,28 @@ def admit(self, step: BoundedKVStep, ctx: StepContext) -> AdmitOutcome: message=f"bounded KV is full: {self.config.max_slots} slots", pages_short=1, label=seg.label, request_id=seg.request_id)) req = self._requests[seg.request_id] = _Request(self._free.pop()) + taken[seg.request_id] = req.slot if pos.written != req.written: return AdmitOutcome(ok=False, ready=False, reason=AdmitRuntimeError( f"{seg.request_id!r} declares {pos.written} written tokens; " f"{req.written} were committed", request_ids=frozenset({seg.request_id}))) return ADMIT_OK + def rollback_admit(self, step: BoundedKVStep, ctx: StepContext) -> None: + """Free the slots this step's admit took fresh; a slot a request already + held stays with it.""" + del step + if ctx.capture or not self._admit_taken: + return + with self._lock: + taken = self._admit_taken + for rid in [r for r in taken if ctx.is_real_row(r)]: + slot = taken.pop(rid) + req = self._requests.get(rid) + if req is not None and req.slot == slot: + del self._requests[rid] + self._free.append(slot) + def plan(self, step: BoundedKVStep, ctx: StepContext) -> BoundedPlan: c = self.config values: list[int] = [] diff --git a/test/modular/test_bounded_kv.py b/test/modular/test_bounded_kv.py index 5aa69549a7..08a3798c20 100644 --- a/test/modular/test_bounded_kv.py +++ b/test/modular/test_bounded_kv.py @@ -404,6 +404,28 @@ def test_admit_refusal_names_the_request(): assert not outcome.ok and outcome.reason.request_ids == frozenset({2}) +def test_rollback_admit_frees_only_fresh_slots(): + """A refused step gives back the slots its admit took, not one a request + already held.""" + from mstar.engine.resources.kv.bounded.config import StreamPosition + + mgr = _manager(WINDOW, True, torch.float32, "cpu", heads=1, head_dim=4, slots=3) + first, ctx = _step([1], [0], [1]) + assert mgr.admit(first, ctx).ok + mgr.plan(first, ctx) + mgr.commit(first, ctx) + # 2 takes a slot, then 3 (a fresh request declaring written tokens) is refused + bad, ctx = _step([1, 2, 3], [1, 0, 0], [1, 1, 1]) + bad = type(bad)(segments=bad.segments, positions={**bad.positions, 3: StreamPosition(0, 5)}) + assert not mgr.admit(bad, ctx).ok + mgr.rollback_admit(bad, ctx) + # both slots are free again, and 1 kept its slot and count + fresh, ctx = _step([4, 5], [0, 0], [1, 1]) + assert mgr.admit(fresh, ctx).ok + again, ctx = _step([1], [1], [1]) + assert mgr.admit(again, ctx).ok + + def test_per_query_window_config(): from mstar.engine.resources.kv.bounded.config import BoundedKVConfig From 710b188703fcc339ad6341aaf41ac7b9d3db79c2 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:58:51 -0700 Subject: [PATCH 52/69] Bounded KV: refuse arguments the kernel would misread attend and bounded_attention now raise ValueError on host facts alone: q, k, v in a dtype other than the cache's; a step whose rows read a source or row_source entry when attend was not given one (the kernel read slot 0 instead); a rel_bias band narrower than the step's largest retained count plus 2T - 1 (per step) or window + 1 (per query) distances, which the kernel read past; and a head_dim the kernel cannot tile (not a power of two >= 16). plan records the facts it already computes for every row. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/kv/bounded/kernels.py | 7 +++- mstar/engine/resources/kv/bounded/manager.py | 34 ++++++++++++++-- test/modular/test_bounded_kv.py | 42 ++++++++++++++++++++ 3 files changed, 78 insertions(+), 5 deletions(-) diff --git a/mstar/engine/resources/kv/bounded/kernels.py b/mstar/engine/resources/kv/bounded/kernels.py index 4558339bda..35b83d016c 100644 --- a/mstar/engine/resources/kv/bounded/kernels.py +++ b/mstar/engine/resources/kv/bounded/kernels.py @@ -218,6 +218,10 @@ def bounded_attention( ``R = window + 1``. Returns ``[B * rows, T, H, D]``.""" n, h, t, d = q.shape rows = cache.shape[1] + if d < 16 or d & (d - 1): + raise ValueError(f"the kernel tiles head_dim, so it must be a power of two >= 16, got {d}") + if not q.dtype == k.dtype == v.dtype == cache.dtype: + raise ValueError(f"the cache holds {cache.dtype}; q, k, v are {q.dtype}, {k.dtype}, {v.dtype}") assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() assert cache.stride(-1) == 1 and table.shape[-1] == ROW_INTS has_row_source = row_source is not None @@ -237,7 +241,8 @@ def bounded_attention( has_bias = rel_bias is not None if has_bias: assert rel_bias.shape[:3] == (n, h, t) and rel_bias.stride(-1) == 1 - assert window is None or rel_bias.shape[3] > window, "a per-query band holds window + 1 distances" + if window is not None and rel_bias.shape[3] <= window: + raise ValueError(f"a per-query band holds window + 1 = {window + 1} distances, got {rel_bias.shape[3]}") s_bias = (rel_bias.stride(0), rel_bias.stride(1), rel_bias.stride(2)) else: rel_bias, s_bias = q, (0, 0, 0) diff --git a/mstar/engine/resources/kv/bounded/manager.py b/mstar/engine/resources/kv/bounded/manager.py index d446f11748..fd7eaaf08b 100644 --- a/mstar/engine/resources/kv/bounded/manager.py +++ b/mstar/engine/resources/kv/bounded/manager.py @@ -34,6 +34,10 @@ class BoundedPlan(NamedTuple): table: torch.Tensor # [>= rows, ROW_INTS] int32 # (slot, layout, source) per row, padding included; what the table holds rows: tuple[tuple[int, RowLayout, int], ...] + # host facts ``attend`` checks its arguments against + max_retained: int = 0 + reads_source: bool = False + reads_row_source: bool = False def retained(self) -> torch.Tensor: """Each row's retained keys before its step, ``[rows]`` int32 on the device.""" @@ -115,6 +119,7 @@ def attend( slots [>= B])``, e.g. another resource's per-request entries.""" plan = plan or self.current cache = self._cache[layer_idx] + self._check_attend(q, k, v, cache, source, plan, rel_bias, row_source) if not q.is_cuda: return self._attend_torch(q, k, v, cache, source, plan, rel_bias, row_source) from mstar.engine.resources.kv.bounded.kernels import bounded_attention, bounded_store @@ -125,6 +130,19 @@ def attend( bounded_store(k, v, cache, plan.table, self.config.reverse_step_order) return out + def _check_attend(self, q, k, v, cache, source, plan: BoundedPlan, rel_bias, row_source) -> None: + if not q.dtype == k.dtype == v.dtype == cache.dtype: + raise ValueError(f"bounded KV holds {cache.dtype}; q, k, v are {q.dtype}, {k.dtype}, {v.dtype}") + if plan.reads_source and source is None: + raise ValueError("a row of this step reads its source prefix, and attend was given no source") + if plan.reads_row_source and row_source is None: + raise ValueError("a row of this step reads its row_source entry, and attend was given no row_source") + if rel_bias is not None: + t, window = q.shape[2], self.config.query_window + need = plan.max_retained + 2 * t - 1 if window is None else window + 1 + if rel_bias.shape[-1] < need: + raise ValueError(f"rel_bias holds {rel_bias.shape[-1]} distances; this step needs {need}") + def _attend_torch( self, q, k, v, cache, source, plan: BoundedPlan, rel_bias=None, row_source=None, ) -> torch.Tensor: @@ -250,6 +268,7 @@ def plan(self, step: BoundedKVStep, ctx: StepContext) -> BoundedPlan: c = self.config values: list[int] = [] rows = [] + max_retained, reads_source, reads_row_source = 0, False, False for seg in step.segments or (): req = None if ctx.capture or ctx.is_padding_row(seg.request_id) else self._requests.get(seg.request_id) if req is None: @@ -257,13 +276,20 @@ def plan(self, step: BoundedKVStep, ctx: StepContext) -> BoundedPlan: rows.append((0, EMPTY_ROW, 0)) else: pos = step.positions[seg.request_id] - rows.append((req.slot, row_layout(pos.source_len, pos.written, seg.span, - c.retention_for(pos.source_len), c.sink_capacity, - c.reverse_step_order), pos.source)) + layout = row_layout(pos.source_len, pos.written, seg.span, c.retention_for(pos.source_len), + c.sink_capacity, c.reverse_step_order) + rows.append((req.slot, layout, pos.source)) + reads = layout.reads + max_retained = max(max_retained, sum(count for _, count in reads)) + if reads[0][1] or reads[2][1]: # the source's ranges (READ_FROM_SOURCE) + if pos.source < 0: + reads_row_source = True + else: + reads_source = True values += flatten_row(*rows[-1]) table = self._table(ctx, len(rows)) self._stager.copy_(table, values, pad_value=0) - self._current = BoundedPlan(table, tuple(rows)) + self._current = BoundedPlan(table, tuple(rows), max_retained, reads_source, reads_row_source) return self._current def commit(self, step: BoundedKVStep, ctx: StepContext) -> None: diff --git a/test/modular/test_bounded_kv.py b/test/modular/test_bounded_kv.py index 08a3798c20..440f457efa 100644 --- a/test/modular/test_bounded_kv.py +++ b/test/modular/test_bounded_kv.py @@ -426,6 +426,48 @@ def test_rollback_admit_frees_only_fresh_slots(): assert mgr.admit(again, ctx).ok +def test_attend_refuses_misuse(): + """Arguments the kernel would read out of bounds or reinterpret are refused + on the host, before a launch.""" + from mstar.engine.resources.kv.bounded.config import BoundedKVConfig, BoundedKVStep, StreamPosition + from mstar.engine.resources.kv.bounded.manager import BoundedKVManager + from mstar.engine.resources.step import Segment, StepContext + + h, d, t = 1, 16, 2 + config = BoundedKVConfig(num_layers=1, num_heads=h, head_dim=d, rows_per_request=1, min_source_len=10, + max_source_len=10, retention=SinkWindow(4, 8), max_slots=2) + mgr = BoundedKVManager(config, torch.device("cpu")) + # row 1 starts from a shared source, row 2 from its own row_source entry + step = BoundedKVStep(segments=(Segment(1, "main", t), Segment(2, "main", t)), + positions={1: StreamPosition(10, 0, 0), 2: StreamPosition(10, 0, -1)}) + ctx = StepContext(request_ids=(1, 2), graph_walk="w", slot=0, capture=False) + assert mgr.admit(step, ctx).ok + mgr.plan(step, ctx) + q, k, v = (torch.randn(2, h, t, d) for _ in range(3)) + source = torch.randn(1, 1, h, 10, 2 * d) + row_source = (torch.randn(2, 1, h, 10, 2 * d), torch.tensor([1, 0], dtype=torch.int32)) + # 10 retained (4 sink + 6 window, all source) + 2T - 1 + band = torch.zeros(2, h, t, 10 + 2 * t - 1) + mgr.attend(0, q, k, v, source, rel_bias=band, row_source=row_source) + for kwargs in (dict(row_source=row_source), dict(source=source), + dict(source=source, row_source=row_source, rel_bias=band[..., :-1])): + with pytest.raises(ValueError): + mgr.attend(0, q, k, v, **kwargs) + with pytest.raises(ValueError): + mgr.attend(0, q.double(), k, v, source, row_source=row_source) + + +def test_kernel_refuses_a_head_dim_it_cannot_tile(): + pytest.importorskip("triton") + from mstar.engine.resources.kv.bounded.kernels import bounded_attention + from mstar.engine.resources.kv.bounded.layout import ROW_INTS + + q = torch.randn(1, 1, 1, 24) + cache = torch.zeros(1, 1, 1, 4, 48) + with pytest.raises(ValueError): + bounded_attention(q, q, q, cache, None, torch.zeros(1, ROW_INTS, dtype=torch.int32)) + + def test_per_query_window_config(): from mstar.engine.resources.kv.bounded.config import BoundedKVConfig From 0abfcf432f24c9190de83987a1a7160756afa2f2 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 17:00:28 -0700 Subject: [PATCH 53/69] Bounded KV: say the online softmax's -inf guard is defensive; undangle a docs path The guard comment and the stream test's docstring read as if a test reached a fully masked stored query. None does while BLOCK_M <= BLOCK_N, since the own-key loop runs first and its first tile holds one of each query's keys; the guard only keeps masked-off queries' unstored accumulators finite. The docs named token2wav's file, which is not on main. Co-Authored-By: Claude Opus 5.5 --- docs/adding_models.rst | 5 ++--- mstar/engine/resources/kv/bounded/kernels.py | 6 ++++-- test/modular/test_bounded_kv.py | 5 ++--- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 4fe7b62ebd..da9fe8c8e2 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -702,9 +702,8 @@ writes none), unmasked, and its writes land after. Pass ``source`` (``[sources, rows, H, L, 2 * head_dim]``, keys then values) when requests start from shared prefixes that the slots should not copy; each row reads the one its ``StreamPosition.source`` names, so rows of different prefixes share a batch and a graph. ``rel_bias`` adds a term that depends only on -the query-key distance, for Transformer-XL-style relative positions. MiniCPM-o's token2wav -(``CACHE_FAMILIES`` in ``mstar/model/minicpm_o/components/token2wav.py``) declares three: -its two conformer caches and its DiT's. +the query-key distance, for Transformer-XL-style relative positions. A streaming vocoder +like MiniCPM-o's token2wav fits three of them: its two conformer caches and its DiT's. Two flags cover retention rules that are not quite a token FIFO. ``SinkWindow(window_includes_step=True)`` counts the step's own tokens toward the window. diff --git a/mstar/engine/resources/kv/bounded/kernels.py b/mstar/engine/resources/kv/bounded/kernels.py index 35b83d016c..0e89595ae8 100644 --- a/mstar/engine/resources/kv/bounded/kernels.py +++ b/mstar/engine/resources/kv/bounded/kernels.py @@ -115,8 +115,10 @@ def _bounded_attention_kernel( rel = offs_m[:, None] - offs_n[None, :] + T - 1 s = _scores(q, k, bias_rows, rel, valid, sm_scale, HAS_BIAS, PRECISION) m_new = tl.maximum(m_i, tl.max(s, axis=1)) - # a query whose keys so far are all masked keeps m = -inf, and - # exp(-inf - -inf) is NaN: offset by 0 instead, so alpha and p are 0 + # Defensive: a query whose keys so far are all masked keeps m = -inf, and + # exp(-inf - -inf) is NaN, so offset by 0 instead. No stored query gets + # here while BLOCK_M <= BLOCK_N: this loop runs first, and its first tile + # holds one of each query's keys. m_off = tl.where(m_new == float("-inf"), 0.0, m_new) alpha = tl.exp(m_i - m_off) p = tl.exp(s - m_off[:, None]) diff --git a/test/modular/test_bounded_kv.py b/test/modular/test_bounded_kv.py index 440f457efa..a08a34bc41 100644 --- a/test/modular/test_bounded_kv.py +++ b/test/modular/test_bounded_kv.py @@ -350,9 +350,8 @@ def _check_stream(mgr, schedule, ieee, tol): ]) def test_stream_matches_brute_force(device, dtype, per_query): """Rows at different positions in one batch, over one-token steps and steps - longer than 2 x BLOCK_N (where a per-query window leaves whole key tiles - masked for some queries), rows padded past their span, and rows that only - read (span 0), at head_dim 128.""" + longer than 2 x BLOCK_N, rows padded past their span, and rows that only read + (span 0), at head_dim 128.""" torch.manual_seed(0) mgr = _manager(WINDOW, per_query, dtype, device, heads=2, head_dim=128) schedule = [(1, [1, 1, 0]), (5, [5, 2, 0]), (150, [150, 0, 150]), (3, [3, 3, 1]), From bb0fb52f33fcb93b333183540a73967eb87151e5 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 17:00:29 -0700 Subject: [PATCH 54/69] Bounded KV: only CPU tensors take the torch path attend ran its per-row torch reference for anything not CUDA, including an accelerator Triton targets; that path also syncs (row_source slots via tolist). It is the CPU fallback, so it is taken on CPU only. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/kv/bounded/manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mstar/engine/resources/kv/bounded/manager.py b/mstar/engine/resources/kv/bounded/manager.py index fd7eaaf08b..c9fd656515 100644 --- a/mstar/engine/resources/kv/bounded/manager.py +++ b/mstar/engine/resources/kv/bounded/manager.py @@ -120,7 +120,7 @@ def attend( plan = plan or self.current cache = self._cache[layer_idx] self._check_attend(q, k, v, cache, source, plan, rel_bias, row_source) - if not q.is_cuda: + if q.device.type == "cpu": return self._attend_torch(q, k, v, cache, source, plan, rel_bias, row_source) from mstar.engine.resources.kv.bounded.kernels import bounded_attention, bounded_store From ff59d3277316b08a57469641f468cf28eb6c7378 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 17:00:29 -0700 Subject: [PATCH 55/69] Bounded KV: resolve dtype=None without writing it into the config build replaced the shared config's dtype=None with the engine's KV dtype, so the model's spec no longer equalled one built afresh from the same declaration. The manager now takes the engine's dtype as a fallback and the config stays as declared. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/kv/bounded/manager.py | 11 ++++++----- test/modular/test_bounded_kv.py | 2 ++ 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/mstar/engine/resources/kv/bounded/manager.py b/mstar/engine/resources/kv/bounded/manager.py index c9fd656515..2bb9689ff1 100644 --- a/mstar/engine/resources/kv/bounded/manager.py +++ b/mstar/engine/resources/kv/bounded/manager.py @@ -58,18 +58,19 @@ class BoundedKVManager(Resource): def build(cls, spec: BoundedKVSpec, info: EngineResourceInfo) -> "BoundedKVManager": if info.joint_comm_group is not None and info.joint_comm_group.world_size > 1: raise NotImplementedError("bounded KV does not shard across ranks yet") - if spec.config.dtype is None: - spec.config.dtype = info.kv_dtype if info.kv_dtype is not None else torch.get_default_dtype() - return cls(spec.config, info.device) + # resolved here, not written back: the config is the model's, shared by its specs + return cls(spec.config, info.device, info.kv_dtype) - def __init__(self, config: BoundedKVConfig, device: torch.device): + def __init__(self, config: BoundedKVConfig, device: torch.device, kv_dtype: torch.dtype | None = None): + """``kv_dtype`` (the engine's) stands in for ``config.dtype=None``.""" self.config = config self._device = torch.device(device) c = config + dtype = c.dtype or kv_dtype or torch.get_default_dtype() # layer-major so one layer's slots are a contiguous [slots, rows, H, cap, 2D] self._cache = torch.zeros( c.num_layers, c.max_slots, c.rows_per_request, c.num_heads, max(c.slot_tokens, 1), - 2 * c.head_dim, dtype=c.dtype, device=self._device, + 2 * c.head_dim, dtype=dtype, device=self._device, ) self._requests: dict[str, _Request] = {} self._free = list(reversed(range(c.max_slots))) diff --git a/test/modular/test_bounded_kv.py b/test/modular/test_bounded_kv.py index a08a34bc41..f19f39a74e 100644 --- a/test/modular/test_bounded_kv.py +++ b/test/modular/test_bounded_kv.py @@ -531,5 +531,7 @@ def test_dtype_none_is_the_engines_kv_dtype(): spec = BoundedKVSpec(resource_key="kv", nodes={"node"}, config=config) mgr = BoundedKVManager.build(spec, EngineResourceInfo(device=torch.device("cpu"), kv_dtype=torch.bfloat16)) assert mgr.layer(0).dtype == torch.bfloat16 + # left as declared, so the spec still equals one built from the same config + assert config.dtype is None # the engine sizes graph tables through main's hook name assert BoundedKVManager.build_accelerator_graph_buffers is not Resource.build_accelerator_graph_buffers From 521566aeaf3b797328fa7f0b1ddfe5ae7b800f4e Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 16:23:12 -0700 Subject: [PATCH 56/69] nemotron_duplex: the encoder's attention window on bounded KV The streaming encoder kept each layer's last 70 keys and values as opaque blocks of its recurrent-state pool, gathered them, masked and scored them in model code, and scattered the window back. They now live in a bounded KV (ENCODER_KV: 24 layers x 8 heads x 128, a per-query window of 70 frames, one slot per session, in the engine's compute dtype), which also runs the attention: _attend computes q, k, v and the relative-position band of distances 0..70 (bd_all, unscaled; the kernel scales by d^-0.5) and calls kv.attend. Declarations: a step's BoundedKVStep places each call at its host StreamState.frames and gives it the frames it adds (n, counted on the host in prepare_inputs by ConformerStreamer.new_frames, which counts as `front` does; the forward asserts the two agree). A steady call that STEP_REGION serves declares span 0 in the outer step, since the region writes its frame; the forward sends exactly those calls to the region. The regions' declarations give each known call one frame at its position and leave capture and replay padding rows (no stream state) unplaced. validate_config_yaml counts encoder_kv's slots against max_concurrent_requests. Retention is a constant SinkWindow with query_window set to the same 70, so a declaration from the defaults equals one from the checkpoint by value (_refuse_resources_sized_otherwise compares specs). Deleted: the "k"/"v" blocks of encoder_state_blocks, _mask and its dist table, the window gather / scatter and kv_idx, the regions' `frames` static input, and `frames` in EncoderInputs and the steady_step / layers signatures (the conv tails keep n). Tests: the encoder parity tests drive the KV's torch path in fp32 at the existing 1e-5, with a step mixing the whole-step region and an eager call whose chunk adds more frames than the window; the resource set, dims (checked against the checkpoint's config.json: 24 layers, 8 heads, d_model 1024, att_context_size [70, 0]), dtype, session cap and checkpoint check cover ENCODER_KV. Co-Authored-By: Claude Opus 5.5 --- configs/nemotron_duplex.yaml | 8 +- configs/nemotron_duplex_disagg.yaml | 8 +- docs/adding_models.rst | 4 +- docs/models.rst | 14 +- .../nemotron_duplex/components/conformer.py | 128 +++++++++--------- mstar/model/nemotron_duplex/config.py | 3 +- .../nemotron_duplex/nemotron_duplex_model.py | 34 +++-- mstar/model/nemotron_duplex/submodules.py | 125 +++++++++++------ test/modular/test_nemotron_duplex_model.py | 59 +++++--- .../test_nemotron_duplex_submodules.py | 88 ++++++++++-- 10 files changed, 308 insertions(+), 163 deletions(-) diff --git a/configs/nemotron_duplex.yaml b/configs/nemotron_duplex.yaml index 80da7e1e8d..68e2d44855 100644 --- a/configs/nemotron_duplex.yaml +++ b/configs/nemotron_duplex.yaml @@ -10,10 +10,10 @@ model: "nemotron_duplex" max_seq_len: 12288 # config: voicechat streaming_max_len / vllm_max_model_len # Sessions at once; the overflow queues. Each session holds a slot in -# resources mamba_state, encoder_state and encoder_front 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 sink) and the KV pools' max_num_pages -# with it. +# 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 diff --git a/configs/nemotron_duplex_disagg.yaml b/configs/nemotron_duplex_disagg.yaml index 7eb234375c..00ec371c17 100644 --- a/configs/nemotron_duplex_disagg.yaml +++ b/configs/nemotron_duplex_disagg.yaml @@ -15,10 +15,10 @@ model: "nemotron_duplex" max_seq_len: 12288 # Sessions at once; the overflow queues. Each session holds a slot in -# resources mamba_state, encoder_state and encoder_front 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 sink) and the KV pools' max_num_pages -# with it. +# 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 diff --git a/docs/adding_models.rst b/docs/adding_models.rst index da9fe8c8e2..19b74d8711 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -721,7 +721,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/models.rst b/docs/models.rst index b2c5bfbb92..67b58917ed 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -443,11 +443,12 @@ Nemotron VoiceChat (``nemotron_duplex``) notes 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 its pages in ``nano_kv`` and + 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 five with it. The server + 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 @@ -530,9 +531,12 @@ 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 two recurrent-state pools, a slot per - session each: its layer state (attention keys / values and conv inputs, - ``encoder_state``) and its newest audio and mel frames (``encoder_front``). A +- 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 diff --git a/mstar/model/nemotron_duplex/components/conformer.py b/mstar/model/nemotron_duplex/components/conformer.py index acb668633c..a81a2100fa 100644 --- a/mstar/model/nemotron_duplex/components/conformer.py +++ b/mstar/model/nemotron_duplex/components/conformer.py @@ -25,12 +25,14 @@ import torch.nn.functional as F from torch import nn +from mstar.engine.resources.kv.bounded import BoundedKVConfig, SinkWindow 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 RecurrentAddressing, RecurrentStatePool # Pure numerical constants (not architecture dims / not in config.json). @@ -342,7 +344,7 @@ class StreamState: 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 layers' state) + self.frames = 0 # encoder frames emitted (= the call's position in the bounded KV) def front_state_blocks(config: ConformerSTTConfig) -> dict[str, RecurrentBlockConfig]: @@ -360,18 +362,26 @@ def front_state_blocks(config: ConformerSTTConfig) -> dict[str, RecurrentBlockCo def encoder_state_blocks(config: ConformerSTTConfig, dtype: torch.dtype | None = None) -> dict[str, RecurrentBlockConfig]: - """A call's layer state, per layer: the last ``att_context`` keys and values - of its attention and the last ``kernel - 1`` inputs of its depthwise conv. + """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.""" - ctx = config.att_context_size[0] return { - "k": RecurrentBlockConfig(shape=(config.num_heads, ctx, config.head_dim), dtype=dtype), - "v": RecurrentBlockConfig(shape=(config.num_heads, ctx, config.head_dim), dtype=dtype), "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: """Mel frames whose 512-sample window is complete, for a call of more than ``n_fft/2`` samples: frame m spans ``[m * hop - n_fft/2, m * hop + n_fft/2)``. @@ -400,9 +410,8 @@ def device_counts(values, device: torch.device) -> torch.Tensor: 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 frames before the step and new ones.""" + most any call got, and per call its new ones.""" x: torch.Tensor - frames: list[int] n: list[int] @@ -418,12 +427,15 @@ class ConformerStreamer: 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. - - A call's state is its host counters (``StreamState``) and two pool slots: - the front end's (``front_state_blocks``) and the layers' - (``encoder_state_blocks``). ``steady_step`` is a whole step for calls that + - 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. """ @@ -491,9 +503,8 @@ def is_steady(self, st: StreamState, chunk: torch.Tensor | None, final: bool) -> 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, frames: torch.Tensor, front: "RecurrentStatePool", - front_rows: "RecurrentAddressing", pool: "RecurrentStatePool", - rows: "RecurrentAddressing") -> torch.Tensor: + def steady_step(self, chunk: torch.Tensor, front: "RecurrentStatePool", front_rows: "RecurrentAddressing", + kv: "BoundedKVManager", pool: "RecurrentStatePool", rows: "RecurrentAddressing") -> 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 @@ -517,7 +528,7 @@ def steady_step(self, chunk: torch.Tensor, frames: torch.Tensor, front: "Recurre front.scatter_("sig", front_rows, y) front.scatter_("last", front_rows, x[:, -1:]) front.scatter_("mel", front_rows, torch.cat([mel_tail, mel_new], dim=1)[:, -2 * f:]) - return self.layers(new, pool, rows, frames, torch.ones_like(frames)) + return self.layers(new, kv, pool, rows, chunk.new_ones(chunk.shape[0], dtype=torch.long)) def front(self, sts: list[StreamState], chunks: list[torch.Tensor | None], final: list[bool], front: "RecurrentStatePool", front_rows: "RecurrentAddressing") -> EncoderInputs | None: @@ -534,7 +545,6 @@ def front(self, sts: list[StreamState], chunks: list[torch.Tensor | None], final mel_tail = front.gather("mel", front_rows) new_sig, new_last, new_mel = list(sig), list(last), list(mel_tail) n = [0] * len(sts) - frames_before = [st.frames for st in sts] rows, lo_hi, mel_rows = [], [], [] for i in busy: st, chunk = sts[i], chunks[i] @@ -592,7 +602,17 @@ def front(self, sts: list[StreamState], chunks: list[torch.Tensor | None], final 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, frames=frames_before, n=n) + 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 as ``front`` counts them, but on the + host before the step runs, so its declaration can name them.""" + 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: @@ -602,42 +622,28 @@ def advance(sts: list[StreamState], enc: EncoderInputs) -> None: # -- layers -------------------------------------------------------------- - def _mask(self, frames: torch.Tensor, n: torch.Tensor, t: int) -> 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 first ``n`` real); the same for every layer.""" - ctx, device = self._ctx, frames.device - 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) - frames, n = frames.view(-1, 1, 1), n.view(-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 >= n) - 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]: @@ -651,33 +657,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, pool: "RecurrentStatePool", rows: "RecurrentAddressing", - frames: torch.Tensor, n: torch.Tensor) -> torch.Tensor: - """The layer stack over ``x (B, t, D)`` and the projection, reading and - writing each call's state in ``pool`` at ``rows``. ``frames`` / ``n`` are - device ``(B,)`` counts: frames before the step, new frames (``n <= t``; a - call with none keeps its state as it was).""" + def layers(self, x: torch.Tensor, kv: "BoundedKVManager", pool: "RecurrentStatePool", + rows: "RecurrentAddressing", 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 ``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 - ctx, keep = self._ctx, self.c.conv_kernel_size - 1 - mask, dist = self._mask(frames, n, t) - # each call keeps the newest ctx keys/values and kernel-1 conv inputs: - # positions n .. n + keep of each concatenation - shift = n.view(b, 1, 1, 1) - kv_idx = (shift + torch.arange(ctx, device=x.device).view(1, 1, ctx, 1)).expand( - b, self.c.num_heads, ctx, self.c.head_dim) - conv_idx = (shift.view(b, 1, 1) + torch.arange(keep, device=x.device).view(1, 1, keep)).expand( + 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): - k_cache, v_cache = pool.gather("k", rows, li), pool.gather("v", rows, li) conv_cache = pool.gather("conv", rows, 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_cache, v_cache, mask, dist) - r = r + a + 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) - pool.scatter_("k", rows, k.gather(2, kv_idx).to(k_cache.dtype), li) - pool.scatter_("v", rows, v.gather(2, kv_idx).to(v_cache.dtype), li) pool.scatter_("conv", rows, tails.gather(2, conv_idx).to(conv_cache.dtype), li) return self.p.proj(x) diff --git a/mstar/model/nemotron_duplex/config.py b/mstar/model/nemotron_duplex/config.py index 6f028d6fb6..41caa01065 100644 --- a/mstar/model/nemotron_duplex/config.py +++ b/mstar/model/nemotron_duplex/config.py @@ -30,8 +30,9 @@ 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 layer state +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``). diff --git a/mstar/model/nemotron_duplex/nemotron_duplex_model.py b/mstar/model/nemotron_duplex/nemotron_duplex_model.py index ca574f5e26..a64d8ee54d 100644 --- a/mstar/model/nemotron_duplex/nemotron_duplex_model.py +++ b/mstar/model/nemotron_duplex/nemotron_duplex_model.py @@ -32,8 +32,8 @@ 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 encoder's two recurrent-state pools (each call's layer state, and its audio / -mel tails), the nano's paged KV cache, attention, recurrent-state pool (Mamba-2 +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). @@ -67,16 +67,22 @@ 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, Model from mstar.model.multimodal import PromptPart -from mstar.model.nemotron_duplex.components.conformer import encoder_state_blocks, front_state_blocks +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, @@ -410,14 +416,15 @@ 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. conformer_encoder: two recurrent-state pools, each - # call's layer state and its front end's audio / mel tails (one slot per - # session each). The codec declares 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 (a slot each in mamba_state, ~137 MB in fp32 for - # Nemotron-H, and in encoder_state / encoder_front; pages in nano_kv and - # talker_kv). A deployment raising it raises all five (see + # 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) @@ -471,8 +478,8 @@ 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``) and its nano and talker pages from - its first step to its last, so a session past any of them waits on admit until + ``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. """ @@ -490,6 +497,8 @@ def validate_config_yaml(self, config: dict, config_path: str) -> None: 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 @@ -542,6 +551,11 @@ def get_node_resources(self) -> list[NodeResourceSpec]: 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"}, diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index c8c153e0f7..b08adba029 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -45,12 +45,14 @@ SlotLease, SubmoduleStep, ) +from mstar.engine.resources.kv.bounded import BoundedKVStep, StreamPosition 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, device_counts from mstar.model.nemotron_duplex.config import ( AUDIO_IN, ENCODER_FRONT, + ENCODER_KV, ENCODER_STATE, MAMBA, MAMBA_STATE, @@ -439,9 +441,10 @@ class ConformerEncoderSubmodule(NodeSubmodule): 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), and the last 70 frames of each - layer's attention keys and values plus the conv tails (``ENCODER_STATE``); - its host counters live in its per-request state. + 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 @@ -499,9 +502,15 @@ def prepare_inputs( final = bool(is_final_stream_chunk) return NodeInputs( tensor_inputs={} if wav is None else {"chunk": wav}, - # whether the whole-step region serves the call: decided here, once, - # for the declaration and the forward to share - kwargs={"final": final, "steady": st is not None and self.streamer.is_steady(st, wav, final)}, + kwargs={ + "final": final, + # whether the whole-step region serves 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 + "frames": 0 if st is None else st.frames, + "n": self.streamer.new_frames(st or StreamState(), samples, final), + }, input_seq_len=0 if wav is None else wav.numel(), # samples, for can_batch ) @@ -526,28 +535,56 @@ def max_batch_size(self, graph_walk: str) -> int: return self.CAPTURE_BATCH_SIZES[-1] @staticmethod - def _declare(request_ids, keys) -> SubmoduleStep: + 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}, + 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] = StreamPosition(source_len=0, written=st.frames) + 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 both pools. 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): + # 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. + to_region = self.STEP_REGION in (piecewise_leases or {}) + if to_region and all(inp.kwargs["steady"] for inp in inputs): return None # Otherwise this step reserves and plans every call's slots, 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, and its admit - # can't refuse a call that has no slot yet. - return self._declare(request_ids, (ENCODER_FRONT, ENCODER_STATE)) + # can't refuse a call that has no slot yet. A steady call STEP_REGION + # serves writes its frame there, so here it only reads (span 0); the + # forward sends exactly those calls to the region. + spans = [0 if to_region and inp.kwargs["steady"] else inp.kwargs["n"] for inp in inputs] + positions = {rid: StreamPosition(source_len=0, written=inp.kwargs["frames"]) + 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: @@ -566,22 +603,19 @@ def get_piecewise_accelerator_graph_configs( 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), - "frames": torch.zeros(shape.bs, dtype=torch.long, device=device), - } + 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), - "frames": torch.zeros(shape.bs, dtype=torch.long, 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), + 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, ) @@ -595,14 +629,14 @@ def region(fn, inputs, keys, lease) -> PiecewiseBatchedConfig: def _step_region(self, call: PiecewiseCallInputs) -> dict[str, torch.Tensor]: front, pool = call.resources[ENCODER_FRONT], call.resources[ENCODER_STATE] - s = call.static_inputs - return {"y": self.streamer.steady_step(s["chunk"], s["frames"], front, front.addressing("main"), - pool, pool.addressing("main"))} + return {"y": self.streamer.steady_step(call.static_inputs["chunk"], front, front.addressing("main"), + call.resources[ENCODER_KV], pool, pool.addressing("main"))} 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"], pool, pool.addressing("main"), s["frames"], s["n"])} + return {"y": self.streamer.layers(s["x"], call.resources[ENCODER_KV], pool, pool.addressing("main"), + s["n"])} def _stream_state(self, rid) -> StreamState: state = self.request_state(rid) @@ -642,14 +676,21 @@ def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, ste runners = engine_inputs.piecewise_runners or {} frames: list[torch.Tensor | None] = [None] * len(rids) - # calls one frame of audio further on the frame grid: the whole step as one graph, run last + # calls one frame of audio further on the frame grid: the whole step as one graph, run last. + # Which those are, the declaration settled: it gave them no write here (span 0). step_runner = runners.get(self.STEP_REGION) - steady = [i for i, s in enumerate(steady) if s] if step_runner is not None else [] - if engine_inputs.step is None: - # declare_step left both pools to the region (see there) - assert len(steady) == len(rids), "an encoder step declared nothing but is not all steady" - elif steady and not step_runner.can_run(len(steady)): - steady = [] + kv_step = None if engine_inputs.step is None else engine_inputs.step.get(ENCODER_KV) + if kv_step is None: + # declare_step left every resource to the region (see there) + assert all(steady), "an encoder step declared nothing but is not all steady" + steady = list(range(len(rids))) + else: + steady = [i for i, (s, seg) in enumerate(zip(steady, kv_step.segments, strict=True)) + if s and seg.span == 0] + if steady and (step_runner is None or not step_runner.can_run(len(steady))): + # their frames can't go through the eager path: its plan writes none for them + raise RuntimeError(f"no {self.STEP_REGION} graph for {len(steady)} calls, though the step " + f"leased one for its {len(rids)}") steady_chunks = [chunks[i] for i in steady] for i in steady: chunks[i] = None # the eager pass leaves these calls as they are @@ -661,28 +702,29 @@ def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, ste enc = None if len(steady) < len(rids): enc = streamer.front(states, chunks, final, *self._pool(engine_inputs, ENCODER_FRONT)) + if kv_step is not None: + # prepare_inputs counted each call's new frames on the host for the declaration + assert (enc.n if enc is not None else [0] * len(rids)) == [seg.span for seg in kv_step.segments], ( + "the encoder's front end made other frame counts than its step declared") if enc is not None: runner = runners.get(self.LAYERS_REGION) new = [i for i, n in enumerate(enc.n) if n] if runner is not None and enc.x.shape[1] == 1 and runner.can_run(len(new)): # only the calls with a frame: the region plans its own rows - counts = device_counts([[enc.frames[i] for i in new], [1] * len(new)], device) y = runner.run(static_inputs={"x": torch.stack([enc.x[i] for i in new]), - "frames": counts[0], "n": counts[1]}, + "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, *self._pool(engine_inputs, ENCODER_STATE), device) + y = self._eager_layers(enc, engine_inputs.resources[ENCODER_KV], + *self._pool(engine_inputs, 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), - "frames": device_counts([states[i].frames for i in steady], device)}, - request_ids=[rids[i] for i in steady], - )["y"] + 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 @@ -698,10 +740,9 @@ def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, ste check_stop_buffers={}, ) - def _eager_layers(self, enc, pool, rows, device) -> list[torch.Tensor]: - """The layer stack eagerly over ``enc`` at ``rows``: each call's new frames.""" - counts = device_counts([enc.frames, enc.n], device) - y = self.streamer.layers(enc.x, pool, rows, counts[0], counts[1]) + def _eager_layers(self, enc, kv, pool, rows, 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, rows, 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: diff --git a/test/modular/test_nemotron_duplex_model.py b/test/modular/test_nemotron_duplex_model.py index 8715a5b741..e2536851f9 100644 --- a/test/modular/test_nemotron_duplex_model.py +++ b/test/modular/test_nemotron_duplex_model.py @@ -23,11 +23,13 @@ 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, @@ -73,25 +75,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 two recurrent pools, each call's layer state - and its front end's audio / mel tails, a slot per session each. The codec - owns 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) == {ENCODER_FRONT, ENCODER_STATE, 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} - node = {ENCODER_STATE: {"conformer_encoder"}, ENCODER_FRONT: {"conformer_encoder"}, + 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 enc.config.blocks["k"].shape == (stt.num_heads, stt.att_context_size[0], stt.head_dim) + 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] @@ -150,7 +159,7 @@ def test_duplex_shipped_configs_turn_websocket_deflate_off(): 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, TALKER_KV, NANO_KV]) +@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 @@ -158,7 +167,8 @@ def test_duplex_refuses_a_session_cap_one_per_session_resource_cannot_hold(key): 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}, TALKER_KV: {"max_num_pages": model.talker_kv_pages(128, 48)}, + 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] @@ -182,19 +192,23 @@ def test_duplex_kv_pools_hold_the_session_cap_at_the_calls_it_serves(kwargs): @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 for the next; in another - dtype than the layers compute it in, live frames drift from one-shot ones. - So the pool takes the dtype the engine resolved (``autocast_dtype: null`` - computes in fp32), and no deployment dtype is refused for it.""" + """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") - (spec,) = [s for s in model.get_node_resources() if s.resource_key == ENCODER_STATE] - pool = RecurrentStatePool.build(spec, EngineResourceInfo(device=torch.device("cpu"), kv_dtype=autocast_dtype)) + 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 spec.config.blocks} == {want} + 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): @@ -239,6 +253,19 @@ def test_duplex_refuses_a_checkpoint_its_resources_were_not_sized_for(monkeypatc _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.""" diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index 8e0cd1a7ca..6bb1a2762e 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -13,14 +13,21 @@ 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 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_state_blocks, front_state_blocks +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, @@ -604,8 +611,9 @@ def _small_perception(**featurizer): class _Encoder: - """The encoder node with its state pool, driven as the engine drives an - eager step: declare, admit and plan the pool, preprocess, forward, commit.""" + """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) @@ -617,7 +625,11 @@ def pool(key, num_layers, blocks): 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.resources = {ENCODER_STATE: self.pool, ENCODER_FRONT: self.front} + 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 @@ -648,8 +660,8 @@ def step(self, rids, inputs) -> BatchedModelOutput: def _step(self, rids, inputs) -> BatchedModelOutput: for rid in rids: - self.pool.ingest_request(rid) - self.front.ingest_request(rid) + 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)) @@ -784,11 +796,19 @@ def test_encoder_layer_region_runs_only_the_calls_with_a_frame(): def test_encoder_paths_read_the_plan_current_when_they_use_it(monkeypatch): """Every read and write of a pool's state, eager or in a region, uses the plan the pool holds at that moment, never one held across another plan: - the eager path runs on this step's plan before a region plans its own. A - call with a chunk over a second runs in a step of its own.""" + the eager path runs on this step's plan before a region plans its own. The + KV attends under a plan of exactly the rows it is given, padding included. + 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) - for pool in enc.resources.values(): + 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) + for pool in (enc.pool, enc.front): for name in ("gather", "scatter_"): def checked(block, rows, *args, _pool=pool, _op=getattr(pool, name), **kwargs): plan = _pool.addressing("main").slot_indices @@ -805,17 +825,26 @@ def checked(block, rows, *args, _pool=pool, _op=getattr(pool, name), **kwargs): 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_both_pools_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 pools: the node declares + 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 both pools.""" + every call in each, and in the KV its new frames at its position, save a + call the region serves, which writes its frame there and here only reads.""" 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)] - assert set(enc.declare([1, 2], nxt).keys()) == {ENCODER_FRONT, ENCODER_STATE} # no region captured + 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] + frames = [enc.node.request_state(h).get(enc.node.STATE_KEY).frames for h in (1, 2)] + assert [step.get(ENCODER_KV).positions[h].written for h in (1, 2)] == frames enc.regions = _encoder_regions(enc, ConformerEncoderSubmodule.STEP_REGION) assert enc.declare([1, 2], nxt) is None for at in range(2560, 1280 * 6, 1280): @@ -823,7 +852,34 @@ def test_encoder_an_all_steady_step_leaves_both_pools_to_its_region(): 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)] - assert set(enc.declare([1, 2], mixed).keys()) == {ENCODER_FRONT, ENCODER_STATE} + step = enc.declare([1, 2], mixed) + assert set(step.keys()) == {ENCODER_FRONT, ENCODER_STATE, ENCODER_KV} + assert _kv_spans(step) == [mixed[0].kwargs["n"], 0] + + +def test_encoder_a_step_mixes_the_steady_region_with_eager_chunks_past_the_window(): + """One step holds a call one frame on (the whole-step region), 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) 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 @pytest.mark.parametrize("b_end_rides", [False, True]) @@ -1044,11 +1100,13 @@ def test_encoder_state_is_per_call_and_cleaned_up(): 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 pool in enc.resources.values(): - pool.remove_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 --------------------------------------------- From e716ec2ac4ba55e3485177a62bb51fd6c13db49c Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 17:09:02 -0700 Subject: [PATCH 57/69] nemotron_duplex: a steady call falls back to the eager path again The whole-step region leases the batch's bucket but replays only the steady calls, and when the runner had dropped the capture for that sub-batch the forward raised: the declaration had given those calls span 0 in the KV, so the eager path could not write their frames. Before the KV moved to the bounded resource they fell back to eager. The declaration now gives every call its new frames, and the forward routes. A steady call takes the region only if the region can replay the steady calls and this step's KV plan goes unused, i.e. the other calls' frames, if any, all go through the layer region: the eager layers would write the steady call's frame from padding into the ring entry the region then reads (test_encoder_a_steady_call_joins_the_eager_layers_of_its_step fails with that check removed). Both decisions read the declared frame counts. The cost is that a steady call batched with a call making several frames in one step runs eagerly with it. Co-Authored-By: Claude Opus 5.5 --- mstar/model/nemotron_duplex/submodules.py | 76 ++++++++++--------- .../test_nemotron_duplex_submodules.py | 63 +++++++++++++-- 2 files changed, 96 insertions(+), 43 deletions(-) diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index b08adba029..2e0b839ed5 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -450,7 +450,9 @@ class ConformerEncoderSubmodule(NodeSubmodule): 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). + 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). @@ -519,6 +521,7 @@ def preprocess(self, graph_walk, engine_inputs, inputs): "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: @@ -569,18 +572,15 @@ def declare_step(self, graph_walk, request_ids, inputs, slot_lease=None, piecewi # 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. - to_region = self.STEP_REGION in (piecewise_leases or {}) - if to_region and all(inp.kwargs["steady"] for inp in inputs): + 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 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, and its admit - # can't refuse a call that has no slot yet. A steady call STEP_REGION - # serves writes its frame there, so here it only reads (span 0); the - # forward sends exactly those calls to the region. - spans = [0 if to_region and inp.kwargs["steady"] else inp.kwargs["n"] for inp in inputs] + # 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: StreamPosition(source_len=0, written=inp.kwargs["frames"]) for rid, inp in zip(request_ids, inputs, strict=True)} return self._declare(request_ids, (ENCODER_FRONT, ENCODER_STATE), @@ -667,7 +667,7 @@ def _to_device(chunks: list, device: torch.device) -> list: moved[i] = part return moved - def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, steady=None, **kwargs): + 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 streamer = self.streamer @@ -676,21 +676,29 @@ def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, ste runners = engine_inputs.piecewise_runners or {} frames: list[torch.Tensor | None] = [None] * len(rids) - # calls one frame of audio further on the frame grid: the whole step as one graph, run last. - # Which those are, the declaration settled: it gave them no write here (span 0). - step_runner = runners.get(self.STEP_REGION) - kv_step = None if engine_inputs.step is None else engine_inputs.step.get(ENCODER_KV) - if kv_step is None: - # declare_step left every resource to the region (see there) + # 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: - steady = [i for i, (s, seg) in enumerate(zip(steady, kv_step.segments, strict=True)) - if s and seg.span == 0] - if steady and (step_runner is None or not step_runner.can_run(len(steady))): - # their frames can't go through the eager path: its plan writes none for them - raise RuntimeError(f"no {self.STEP_REGION} graph for {len(steady)} calls, though the step " - f"leased one for its {len(rids)}") + # 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] for i in steady: chunks[i] = None # the eager pass leaves these calls as they are @@ -702,18 +710,16 @@ def forward_batched(self, graph_walk, engine_inputs, chunk=None, final=None, ste enc = None if len(steady) < len(rids): enc = streamer.front(states, chunks, final, *self._pool(engine_inputs, ENCODER_FRONT)) - if kv_step is not None: - # prepare_inputs counted each call's new frames on the host for the declaration - assert (enc.n if enc is not None else [0] * len(rids)) == [seg.span for seg in kv_step.segments], ( - "the encoder's front end made other frame counts than its step declared") + # prepare_inputs counted each call's new frames on the host for the declaration + assert (enc.n if enc is not None else [0] * len(rids)) == [0 if i in steady else k for i, k in enumerate(n)], ( + "the encoder's front end made other frame counts than its step declared") if enc is not None: - runner = runners.get(self.LAYERS_REGION) - new = [i for i, n in enumerate(enc.n) if n] - if runner is not None and enc.x.shape[1] == 1 and runner.can_run(len(new)): + 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 = 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 = 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], diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index 6bb1a2762e..649b0d04ef 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -833,8 +833,7 @@ def test_encoder_an_all_steady_step_leaves_both_pools_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, save a - call the region serves, which writes its frame there and here only reads.""" + 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): @@ -854,14 +853,14 @@ def test_encoder_an_all_steady_step_leaves_both_pools_to_its_region(): 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"], 0] + assert _kv_spans(step) == [mixed[0].kwargs["n"], 1] -def test_encoder_a_step_mixes_the_steady_region_with_eager_chunks_past_the_window(): - """One step holds a call one frame on (the whole-step region), 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) and a call under - a frame a step: each gets its offline frames.""" +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] @@ -882,6 +881,54 @@ def spy_front(*args): 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: []} + 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): + 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]) @pytest.mark.parametrize("a_chunk", [1280, 9000]) @pytest.mark.parametrize("b_len", [16000 * 2, 16000 * 2 + 640, 16000 * 2 + 777]) From b5124e0141c23198a7c7d4c2362fea41fb3aac6c Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 17:09:39 -0700 Subject: [PATCH 58/69] nemotron_duplex: the front end makes the frames prepare_inputs counted ConformerStreamer.front recounted each call's new encoder frames the way new_frames had already counted them on the host for the declaration, and the forward asserted the two agreed. front now takes the declared counts. Co-Authored-By: Claude Opus 5.5 --- .../nemotron_duplex/components/conformer.py | 16 +++++++--------- mstar/model/nemotron_duplex/submodules.py | 8 +++----- 2 files changed, 10 insertions(+), 14 deletions(-) diff --git a/mstar/model/nemotron_duplex/components/conformer.py b/mstar/model/nemotron_duplex/components/conformer.py index a81a2100fa..885bd1f81d 100644 --- a/mstar/model/nemotron_duplex/components/conformer.py +++ b/mstar/model/nemotron_duplex/components/conformer.py @@ -530,11 +530,11 @@ def steady_step(self, chunk: torch.Tensor, front: "RecurrentStatePool", front_ro front.scatter_("mel", front_rows, torch.cat([mel_tail, mel_new], dim=1)[:, -2 * f:]) return self.layers(new, kv, pool, rows, chunk.new_ones(chunk.shape[0], dtype=torch.long)) - def front(self, sts: list[StreamState], chunks: list[torch.Tensor | None], final: list[bool], + def front(self, sts: list[StreamState], chunks: list[torch.Tensor | None], final: list[bool], n: list[int], front: "RecurrentStatePool", front_rows: "RecurrentAddressing") -> EncoderInputs | None: - """Append each call's chunk; the encoder-input frames it made final, or - None when no call has one. A call with no chunk and not at its end is - left as it is.""" + """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)) @@ -544,7 +544,6 @@ def front(self, sts: list[StreamState], chunks: list[torch.Tensor | None], final last, sig = front.gather("last", front_rows), front.gather("sig", front_rows) mel_tail = front.gather("mel", front_rows) new_sig, new_last, new_mel = list(sig), list(last), list(mel_tail) - n = [0] * len(sts) rows, lo_hi, mel_rows = [], [], [] for i in busy: st, chunk = sts[i], chunks[i] @@ -586,8 +585,6 @@ def front(self, sts: list[StreamState], chunks: list[torch.Tensor | None], final 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)) - # also without a new mel frame: the end of the call can make its last frame final - n[i] = max(0, final_encoder_frames(c, st.mels, st.samples, final[i]) - st.frames) 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]))) @@ -606,8 +603,9 @@ def front(self, sts: list[StreamState], chunks: list[torch.Tensor | None], final 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 as ``front`` counts them, but on the - host before the step runs, so its declaration can name them.""" + 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 diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index 2e0b839ed5..5fc918a219 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -700,8 +700,9 @@ def layers_graph(rows) -> bool: 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] = None # the eager pass leaves these calls as they are + 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 @@ -709,10 +710,7 @@ def layers_graph(rows) -> bool: # a step of its own (see can_batch). enc = None if len(steady) < len(rids): - enc = streamer.front(states, chunks, final, *self._pool(engine_inputs, ENCODER_FRONT)) - # prepare_inputs counted each call's new frames on the host for the declaration - assert (enc.n if enc is not None else [0] * len(rids)) == [0 if i in steady else k for i, k in enumerate(n)], ( - "the encoder's front end made other frame counts than its step declared") + enc = streamer.front(states, chunks, final, eager_n, *self._pool(engine_inputs, ENCODER_FRONT)) if enc is not None: new = [i for i, k in enumerate(enc.n) if k] if layers_graph(new): From 68cd0d8cffe12436b8f7bdf4369a35071ff2b950 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 17:10:45 -0700 Subject: [PATCH 59/69] nemotron_duplex: a call's KV position is built once per step The outer declaration built a StreamPosition per call from prepare_inputs' frame count and every region replay built them again from the stream state. StreamState now holds the call's position, rebuilt only when its frames move (frames reads through it), and both declarations hand that object on. Co-Authored-By: Claude Opus 5.5 --- .../nemotron_duplex/components/conformer.py | 19 +++++++++++++++---- mstar/model/nemotron_duplex/submodules.py | 16 ++++++++-------- .../test_nemotron_duplex_submodules.py | 6 ++++-- 3 files changed, 27 insertions(+), 14 deletions(-) diff --git a/mstar/model/nemotron_duplex/components/conformer.py b/mstar/model/nemotron_duplex/components/conformer.py index 885bd1f81d..950281fff0 100644 --- a/mstar/model/nemotron_duplex/components/conformer.py +++ b/mstar/model/nemotron_duplex/components/conformer.py @@ -25,7 +25,7 @@ import torch.nn.functional as F from torch import nn -from mstar.engine.resources.kv.bounded import BoundedKVConfig, SinkWindow +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 @@ -339,12 +339,22 @@ class StreamState: 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", "frames") + __slots__ = ("samples", "mels", "position") def __init__(self): self.samples = 0 # samples of the call so far self.mels = 0 # mel frames computed - self.frames = 0 # encoder frames emitted (= the call's position in the bounded KV) + # 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]: @@ -616,7 +626,8 @@ def new_frames(self, st: StreamState, samples: int, final: bool) -> int: 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): - st.frames += n + if n: + st.frames += n # -- layers -------------------------------------------------------------- diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index 5fc918a219..1ba47e8bd6 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -45,7 +45,7 @@ SlotLease, SubmoduleStep, ) -from mstar.engine.resources.kv.bounded import BoundedKVStep, StreamPosition +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, device_counts @@ -502,16 +502,17 @@ def prepare_inputs( 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": final, - # whether the whole-step region serves the call: decided here, - # once, for the declaration and the forward to share + # 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 - "frames": 0 if st is None else st.frames, - "n": self.streamer.new_frames(st or StreamState(), samples, final), + "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 ) @@ -562,7 +563,7 @@ def _region_kv_step(self, request_ids) -> BoundedKVStep: 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] = StreamPosition(source_len=0, written=st.frames) + 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, @@ -581,8 +582,7 @@ def declare_step(self, graph_walk, request_ids, inputs, slot_lease=None, piecewi # 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: StreamPosition(source_len=0, written=inp.kwargs["frames"]) - for rid, inp in zip(request_ids, inputs, strict=True)} + 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)) diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index 649b0d04ef..b0e0d9cd23 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -842,8 +842,10 @@ def test_encoder_an_all_steady_step_leaves_both_pools_to_its_region(): 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] - frames = [enc.node.request_state(h).get(enc.node.STATE_KEY).frames for h in (1, 2)] - assert [step.get(ENCODER_KV).positions[h].written for h in (1, 2)] == frames + 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): From 0b7b7de9c0d11d2f7f507e9b1a18b529d275a41f Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 17:10:45 -0700 Subject: [PATCH 60/69] nemotron_duplex: name the all-steady declaration test for its three resources Co-Authored-By: Claude Opus 5.5 --- test/modular/test_nemotron_duplex_submodules.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index b0e0d9cd23..400b2b8c48 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -829,7 +829,7 @@ 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_both_pools_to_its_region(): +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 From 0e7fb5cc21b57ea5698c2e514281711b3192cec0 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Wed, 21 Oct 2026 18:26:19 -0700 Subject: [PATCH 61/69] nemotron_duplex: the module docstring lists the encoder's bounded KV Co-Authored-By: Claude Opus 5.5 --- mstar/model/nemotron_duplex/submodules.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/mstar/model/nemotron_duplex/submodules.py b/mstar/model/nemotron_duplex/submodules.py index 1ba47e8bd6..53e52c8d98 100644 --- a/mstar/model/nemotron_duplex/submodules.py +++ b/mstar/model/nemotron_duplex/submodules.py @@ -2,8 +2,9 @@ Nodes, and the resources each declares (``NemotronDuplexModel.get_node_resources``): - conformer_encoder encoder_state, 16 kHz speech -> per-frame LLM embeds (+ RNN-T); per - encoder_front session, its layers' state and its audio / mel tails + 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 @@ -12,9 +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 keeps a call's state in its 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. +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 From f4ac338b989e012e111d6e153ad7ba43888d0b0d Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Tue, 27 Oct 2026 01:15:34 -0700 Subject: [PATCH 62/69] Merge follow-up: the encoder plan test checks only the KV's plan under #891's gather / scatter_ Main's RecurrentStatePool.gather / scatter_ (#891) take no addressing: they read the plan the pool holds when they run, so the test's wrapper, written for the (name, rows) API, broke the call (TypeError) rather than checking anything. Co-Authored-By: Claude Opus 5.5 --- .../modular/test_nemotron_duplex_submodules.py | 18 ++++-------------- 1 file changed, 4 insertions(+), 14 deletions(-) diff --git a/test/modular/test_nemotron_duplex_submodules.py b/test/modular/test_nemotron_duplex_submodules.py index 400b2b8c48..5054832e0b 100644 --- a/test/modular/test_nemotron_duplex_submodules.py +++ b/test/modular/test_nemotron_duplex_submodules.py @@ -794,11 +794,10 @@ def test_encoder_layer_region_runs_only_the_calls_with_a_frame(): def test_encoder_paths_read_the_plan_current_when_they_use_it(monkeypatch): - """Every read and write of a pool's state, eager or in a region, uses the - plan the pool holds at that moment, never one held across another plan: - the eager path runs on this step's plan before a region plans its own. The - KV attends under a plan of exactly the rows it is given, padding included. - A call with a chunk over a second runs in a step of its own.""" + """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 @@ -808,15 +807,6 @@ def attend_checked(layer_idx, q, *args, **kwargs): return attend(layer_idx, q, *args, **kwargs) monkeypatch.setattr(enc.kv, "attend", attend_checked) - for pool in (enc.pool, enc.front): - for name in ("gather", "scatter_"): - def checked(block, rows, *args, _pool=pool, _op=getattr(pool, name), **kwargs): - plan = _pool.addressing("main").slot_indices - assert rows.slot_indices.untyped_storage().data_ptr() == plan.untyped_storage().data_ptr(), ( - "a pool read through a plan it no longer holds") - return _op(block, rows, *args, **kwargs) - - monkeypatch.setattr(pool, name, 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}) From 27925f8cdc1f327eb7ffea22c597358c34ff5386 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Tue, 27 Oct 2026 01:28:15 -0700 Subject: [PATCH 63/69] P3b: nemotron_duplex configs (#595 wording): a live call past the cap gets a 503 Co-Authored-By: Claude Opus 5.5 --- configs/nemotron_duplex.yaml | 3 ++- configs/nemotron_duplex_disagg.yaml | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/configs/nemotron_duplex.yaml b/configs/nemotron_duplex.yaml index 68e2d44855..c49e0b8e24 100644 --- a/configs/nemotron_duplex.yaml +++ b/configs/nemotron_duplex.yaml @@ -9,7 +9,8 @@ model: "nemotron_duplex" max_seq_len: 12288 # config: voicechat streaming_max_len / vllm_max_model_len -# Sessions at once; the overflow queues. Each session holds a slot in +# 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 diff --git a/configs/nemotron_duplex_disagg.yaml b/configs/nemotron_duplex_disagg.yaml index 00ec371c17..4b146a19c7 100644 --- a/configs/nemotron_duplex_disagg.yaml +++ b/configs/nemotron_duplex_disagg.yaml @@ -14,7 +14,8 @@ model: "nemotron_duplex" max_seq_len: 12288 -# Sessions at once; the overflow queues. Each session holds a slot in +# 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 From f0ef4a579a04ccdec4a0d2553fe5d9fff0c5ac71 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Tue, 27 Oct 2026 01:42:22 -0700 Subject: [PATCH 64/69] P3a: the recurrent pool resolves a dtype=None block without writing it back build() wrote the engine's kv_dtype into the model's RecurrentStateConfig, which its specs (and a second engine) share. Resolve it into the pool's own blocks instead, as the bounded KV does. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/recurrent/pool.py | 16 +++++++++------- test/modular/test_recurrent_gather_scatter.py | 17 +++++++++++++++++ 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/mstar/engine/resources/recurrent/pool.py b/mstar/engine/resources/recurrent/pool.py index fb0190c76f..610997f8bc 100644 --- a/mstar/engine/resources/recurrent/pool.py +++ b/mstar/engine/resources/recurrent/pool.py @@ -76,12 +76,14 @@ def build( config = spec.config if info.joint_comm_group is not None: config.shard(info.joint_comm_group.world_size) - for block in config.blocks.values(): - if block.dtype is None: - block.dtype = info.kv_dtype if info.kv_dtype is not None else torch.get_default_dtype() - 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 @@ -93,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.dtype or kv_dtype or torch.get_default_dtype(), device=device, ) for name, block in config.blocks.items() 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 From 51a07e0a47966f42b00f03b3f6fe9b65f11b376f Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Tue, 27 Oct 2026 01:44:10 -0700 Subject: [PATCH 65/69] P3a: size the recurrent pool's startup log from its resolved blocks The config's byte counts size a dtype=None block in torch's default dtype now that the pool no longer writes its resolved dtype back. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/recurrent/pool.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mstar/engine/resources/recurrent/pool.py b/mstar/engine/resources/recurrent/pool.py index 610997f8bc..28de03a558 100644 --- a/mstar/engine/resources/recurrent/pool.py +++ b/mstar/engine/resources/recurrent/pool.py @@ -132,12 +132,14 @@ def __init__( 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()}, ) From f3d3c1a05470ee0a80d9e6c7ba7ed2b515c15bc1 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Tue, 27 Oct 2026 01:59:24 -0700 Subject: [PATCH 66/69] R2.3: linear-attn managers take the state dtype the pool stores The factory read the block's declared dtype from the shared config, which is None for a block in the engine dtype since P3a. RecurrentStatePool. block_dtype is the one rule the pool allocates by, and the factory asks it. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/linear_attn/base.py | 5 +++-- mstar/engine/resources/recurrent/pool.py | 9 ++++++++- test/modular/test_mamba2_resource.py | 13 +++++++++++++ 3 files changed, 24 insertions(+), 3 deletions(-) diff --git a/mstar/engine/resources/linear_attn/base.py b/mstar/engine/resources/linear_attn/base.py index 08822b7653..1bffa0d46e 100644 --- a/mstar/engine/resources/linear_attn/base.py +++ b/mstar/engine/resources/linear_attn/base.py @@ -12,6 +12,7 @@ LinearAttnVariant, ) from mstar.engine.resources.recurrent.config import DeltaNetGeometry, Mamba2Geometry +from mstar.engine.resources.recurrent.pool import RecurrentStatePool logger = logging.getLogger(__name__) @@ -40,7 +41,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=RecurrentStatePool.block_dtype(pool_config.blocks["ssm"], info.kv_dtype), has_sink=not pool_config.disable_sink_slot, device=info.device, ) @@ -55,7 +56,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=RecurrentStatePool.block_dtype(pool_config.blocks["state"], info.kv_dtype), has_sink=not pool_config.disable_sink_slot, device=info.device, ) diff --git a/mstar/engine/resources/recurrent/pool.py b/mstar/engine/resources/recurrent/pool.py index 28de03a558..55a618b6ca 100644 --- a/mstar/engine/resources/recurrent/pool.py +++ b/mstar/engine/resources/recurrent/pool.py @@ -16,6 +16,7 @@ from mstar.engine.resources.base import CGSlotKey, CGSlotSpec, EngineResourceInfo, Resource from mstar.engine.resources.recurrent.config import ( + RecurrentBlockConfig, RecurrentStateConfig, RecurrentStateSpec, RecurrentStep, @@ -95,7 +96,7 @@ def __init__( self._blocks: dict[str, torch.Tensor] = { name: torch.zeros( (config.num_layers, config.max_slots, *block.shape), - dtype=block.dtype or kv_dtype or torch.get_default_dtype(), + dtype=self.block_dtype(block, kv_dtype), device=device, ) for name, block in config.blocks.items() @@ -145,6 +146,12 @@ def __init__( # Storage access + @staticmethod + def block_dtype(block: RecurrentBlockConfig, kv_dtype: torch.dtype | None) -> torch.dtype: + """The dtype the pool stores ``block`` in; a resource planned against + the pool reads it here rather than from the shared config.""" + return block.dtype or kv_dtype or torch.get_default_dtype() + def block(self, name: str, layer_idx: int) -> torch.Tensor: """One layer's contiguous slot-major view of a block: ``[max_slots, *shape]``.""" return self._blocks[name][layer_idx] 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]) From e1d22cc7994700c32e0b40f1f09a75ab799ded5d Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Tue, 27 Oct 2026 02:00:04 -0700 Subject: [PATCH 67/69] R2.4: delete RecurrentBlockConfig's unused byte sizes numel, nbytes, slot_bytes and total_bytes had no caller since the pool sizes its startup log from its own tensors (P3a), and were wrong for a dtype=None block. Their only test restated them. Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/recurrent/config.py | 20 +------------------- test/modular/test_mamba2_geometry.py | 5 ----- 2 files changed, 1 insertion(+), 24 deletions(-) diff --git a/mstar/engine/resources/recurrent/config.py b/mstar/engine/resources/recurrent/config.py index 58a911f38a..72073e7900 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 @@ -25,7 +24,7 @@ class RecurrentBlockConfig: ``shape`` is opaque to the pool. ``shard_dims`` names the axes divided across ranks. A ``dtype`` of None is the engine's compute dtype (what its - KV caches store), settled when the pool is built. + KV caches store), resolved by the pool, which never writes it back. """ shape: tuple[int, ...] @@ -44,14 +43,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 @@ -277,15 +268,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): config: RecurrentStateConfig diff --git a/test/modular/test_mamba2_geometry.py b/test/modular/test_mamba2_geometry.py index 99612584ca..28643499ad 100644 --- a/test/modular/test_mamba2_geometry.py +++ b/test/modular/test_mamba2_geometry.py @@ -8,7 +8,6 @@ from mstar.engine.resources.recurrent import ( DeltaNetGeometry, Mamba2Geometry, - RecurrentStateConfig, ) NEMOTRON_H = Mamba2Geometry(num_heads=128, head_dim=80, state_size=128, n_groups=8, conv_kernel_size=4) @@ -19,10 +18,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(): From afa4364b53a035507858fdc53d3c2abd4ed4d3f4 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Tue, 27 Oct 2026 02:09:16 -0700 Subject: [PATCH 68/69] R3.2: resolve a recurrent block's dtype on RecurrentBlockConfig The rule moves from the static RecurrentStatePool.block_dtype to RecurrentBlockConfig.resolved_dtype, which the pool and the linear-attention builder both call, so linear_attn/base.py no longer imports the pool. The to_blocks state_dtype hints admit None (the engine's dtype). Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/linear_attn/base.py | 5 ++--- mstar/engine/resources/recurrent/config.py | 11 ++++++++--- mstar/engine/resources/recurrent/pool.py | 9 +-------- 3 files changed, 11 insertions(+), 14 deletions(-) diff --git a/mstar/engine/resources/linear_attn/base.py b/mstar/engine/resources/linear_attn/base.py index 1bffa0d46e..3ace685cf0 100644 --- a/mstar/engine/resources/linear_attn/base.py +++ b/mstar/engine/resources/linear_attn/base.py @@ -12,7 +12,6 @@ LinearAttnVariant, ) from mstar.engine.resources.recurrent.config import DeltaNetGeometry, Mamba2Geometry -from mstar.engine.resources.recurrent.pool import RecurrentStatePool logger = logging.getLogger(__name__) @@ -41,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=RecurrentStatePool.block_dtype(pool_config.blocks["ssm"], info.kv_dtype), + state_dtype=pool_config.blocks["ssm"].resolved_dtype(info.kv_dtype), has_sink=not pool_config.disable_sink_slot, device=info.device, ) @@ -56,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=RecurrentStatePool.block_dtype(pool_config.blocks["state"], info.kv_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 72073e7900..c984c818f1 100644 --- a/mstar/engine/resources/recurrent/config.py +++ b/mstar/engine/resources/recurrent/config.py @@ -24,7 +24,7 @@ class RecurrentBlockConfig: ``shape`` is opaque to the pool. ``shard_dims`` names the axes divided across ranks. A ``dtype`` of None is the engine's compute dtype (what its - KV caches store), resolved by the pool, which never writes it back. + KV caches store), resolved by ``resolved_dtype`` and never written back. """ shape: tuple[int, ...] @@ -35,6 +35,11 @@ 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 @@ -83,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.""" @@ -175,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 diff --git a/mstar/engine/resources/recurrent/pool.py b/mstar/engine/resources/recurrent/pool.py index 55a618b6ca..6194c8c00c 100644 --- a/mstar/engine/resources/recurrent/pool.py +++ b/mstar/engine/resources/recurrent/pool.py @@ -16,7 +16,6 @@ from mstar.engine.resources.base import CGSlotKey, CGSlotSpec, EngineResourceInfo, Resource from mstar.engine.resources.recurrent.config import ( - RecurrentBlockConfig, RecurrentStateConfig, RecurrentStateSpec, RecurrentStep, @@ -96,7 +95,7 @@ def __init__( self._blocks: dict[str, torch.Tensor] = { name: torch.zeros( (config.num_layers, config.max_slots, *block.shape), - dtype=self.block_dtype(block, kv_dtype), + dtype=block.resolved_dtype(kv_dtype), device=device, ) for name, block in config.blocks.items() @@ -146,12 +145,6 @@ def __init__( # Storage access - @staticmethod - def block_dtype(block: RecurrentBlockConfig, kv_dtype: torch.dtype | None) -> torch.dtype: - """The dtype the pool stores ``block`` in; a resource planned against - the pool reads it here rather than from the shared config.""" - return block.dtype or kv_dtype or torch.get_default_dtype() - def block(self, name: str, layer_idx: int) -> torch.Tensor: """One layer's contiguous slot-major view of a block: ``[max_slots, *shape]``.""" return self._blocks[name][layer_idx] From 0587065c7c4d37abcee78222a78d5b5776e6a7a8 Mon Sep 17 00:00:00 2001 From: Joseph Tan Date: Tue, 27 Oct 2026 02:09:16 -0700 Subject: [PATCH 69/69] R3.4: two blank lines before RecurrentStateSpec; fold a test import Co-Authored-By: Claude Opus 5.5 --- mstar/engine/resources/recurrent/config.py | 1 + test/modular/test_mamba2_geometry.py | 5 +---- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/mstar/engine/resources/recurrent/config.py b/mstar/engine/resources/recurrent/config.py index c984c818f1..26ef92ec22 100644 --- a/mstar/engine/resources/recurrent/config.py +++ b/mstar/engine/resources/recurrent/config.py @@ -273,6 +273,7 @@ def shard(self, num_shards: int) -> None: for block in self.blocks.values(): block.shard(num_shards) + @dataclass class RecurrentStateSpec(NodeResourceSpec): config: RecurrentStateConfig diff --git a/test/modular/test_mamba2_geometry.py b/test/modular/test_mamba2_geometry.py index 28643499ad..6051c70c6e 100644 --- a/test/modular/test_mamba2_geometry.py +++ b/test/modular/test_mamba2_geometry.py @@ -5,10 +5,7 @@ import pytest import torch -from mstar.engine.resources.recurrent import ( - DeltaNetGeometry, - Mamba2Geometry, -) +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)