Spaces:
Running on Zero
Running on Zero
| """Slot-based, static-state, CUDA-graphed one-frame Mimi decoder. | |
| Arithmetic mirrors serve/codec.py::StreamingMimi (the verified reference): causal | |
| convolution history, transposed-convolution overlap tails (bias removed from the | |
| stored tail), and sliding-window attention with absolute RoPE positions. The | |
| difference is where state lives: every stateful layer owns ONE static tensor with | |
| a row per slot (row ``max_slots`` is a dummy used to pad a batch to its bucket), | |
| and the decoder transformer runs against static per-slot KV ring buffers, so one | |
| decode is a single CUDA-graph replay with no per-seat Python work. | |
| No upstream module or parameter is modified; HF sub-modules are called for their | |
| weights. Global side effect: warmup() sets torch.backends.cudnn.benchmark=False so | |
| convolution algorithm choice is settled before capture. Matmul/cuDNN TF32 flags | |
| are read at capture time and baked into the graphs. | |
| Sliding-window semantics (HF create_sliding_window_causal_mask, reproduced by | |
| StreamingMimi keeping sliding_window-1 cached tokens plus the new ones): a query | |
| at absolute position p sees keys with p - sliding_window < key_pos <= p. | |
| """ | |
| import torch | |
| import torch.nn.functional as F | |
| from transformers.models.mimi.modeling_mimi import (MimiConv1d, MimiConvTranspose1d, MimiResnetBlock, | |
| apply_rotary_pos_emb, repeat_kv) | |
| NEVER = -(1 << 40) # key_pos of an empty ring entry; always outside every window | |
| class _Bucket: | |
| __slots__ = ('graph', 'slots', 'codes', 'out', 'pinned', 'events', 'turn') | |
| class StaticMimi: | |
| def __init__(self, mimi, max_slots, batch_buckets=(1, 2, 4, 8, 16, 24, 32, 48, 64, 80, 96), ring=256): | |
| cfg = mimi.config | |
| if not cfg.use_causal_conv or cfg.pad_mode != 'constant' or cfg.trim_right_ratio != 1: | |
| raise ValueError('Static decoder requires causal constant padding and trim_right_ratio=1') | |
| if mimi.upsample is None: | |
| raise ValueError('Static decoder expects the Mimi upsample layer') | |
| self.m = mimi | |
| self.max_slots = int(max_slots) | |
| self.dummy = self.max_slots | |
| self.batch_buckets = tuple(sorted(set(int(b) for b in batch_buckets))) | |
| self.window = int(cfg.sliding_window) | |
| self.frame_tokens = int(mimi.upsample.conv.stride[0]) | |
| self.ring = int(ring) | |
| if self.ring < self.window + self.frame_tokens: | |
| raise ValueError('ring must hold sliding_window + frame_tokens entries') | |
| parameter = next(mimi.parameters()) | |
| self.device, self.dtype = parameter.device, parameter.dtype | |
| self.codebooks = int(cfg.num_quantizers) | |
| with torch.inference_mode(): | |
| # Convolution history / overlap tails: one flat row per slot, one view per stateful layer. | |
| self._spec = {} | |
| self._plan(mimi.upsample, 'upsample') | |
| for i, layer in enumerate(mimi.decoder.layers): | |
| self._plan(layer, f'decoder.{i}') | |
| total = sum(c * p for c, p in self._spec.values()) | |
| rows = self.max_slots + 1 | |
| self.conv_flat = torch.zeros(rows, total, device=self.device, dtype=self.dtype) | |
| self.state, offset = {}, 0 | |
| for key, (channels, pad) in self._spec.items(): | |
| self.state[key] = self.conv_flat[:, offset:offset + channels * pad].view(rows, channels, pad) | |
| offset += channels * pad | |
| layers = mimi.decoder_transformer.layers | |
| attn = layers[0].self_attn | |
| self.kv_heads, self.head_dim, self.kv_groups = attn.num_key_value_heads, attn.head_dim, attn.num_key_value_groups | |
| # [layer, k/v, slot, heads, ring, head_dim]; keys are stored post-RoPE. | |
| self.kv = torch.zeros(len(layers), 2, rows, self.kv_heads, self.ring, self.head_dim, | |
| device=self.device, dtype=self.dtype) | |
| self.key_pos = torch.full((rows, self.ring), NEVER, device=self.device, dtype=torch.long) | |
| self.tokens = torch.zeros(rows, device=self.device, dtype=torch.long) | |
| self._steps = torch.arange(self.frame_tokens, device=self.device) | |
| rotary = attn.rotary_emb | |
| self._inv_freq = rotary.inv_freq.detach().float().clone() | |
| self._rope_scale = float(rotary.attention_scaling) | |
| self._pool = None | |
| self._buckets = {} | |
| self.samples_per_frame = None | |
| # ---- state layout ------------------------------------------------------------------------- | |
| def _plan(self, module, key): | |
| if isinstance(module, MimiConv1d): | |
| if module.conv.stride != (1,): | |
| raise ValueError('Only stride-one decoder convolutions are supported') | |
| padding = (module.conv.kernel_size[0] - 1) * module.conv.dilation[0] | |
| if padding: | |
| self._spec[key] = (module.conv.in_channels, padding) | |
| elif isinstance(module, MimiConvTranspose1d): | |
| if module.padding_left != 0 or module.conv.dilation != (1,): | |
| raise ValueError('Unsupported transpose-convolution geometry') | |
| if module.padding_right: | |
| self._spec[key] = (module.conv.out_channels, module.padding_right) | |
| elif isinstance(module, MimiResnetBlock): | |
| self._plan(module.shortcut, key + '.shortcut') | |
| for i, child in enumerate(module.block): | |
| self._plan(child, key + f'.block.{i}') | |
| def state_bytes_per_slot(self): | |
| return {'conv': self.conv_flat[0].numel() * self.conv_flat.element_size(), | |
| 'attention_kv': self.kv[:, :, 0].numel() * self.kv.element_size(), | |
| 'attention_index': self.key_pos[0].numel() * 8 + 8} | |
| # ---- graph body --------------------------------------------------------------------------- | |
| def _layer(self, module, x, slots, key): | |
| if isinstance(module, MimiConv1d): | |
| if key in self.state: | |
| state = self.state[key] | |
| joined = torch.cat((state.index_select(0, slots), x), dim=-1) | |
| state.index_copy_(0, slots, joined[..., -state.shape[-1]:]) | |
| x = joined | |
| return module.conv(x) | |
| if isinstance(module, MimiConvTranspose1d): | |
| y = module.conv(x) | |
| produced = x.shape[-1] * module.conv.stride[0] | |
| if key in self.state: | |
| state = self.state[key] | |
| overlap = state.shape[-1] | |
| if overlap > produced: | |
| raise ValueError('Chunk too short for transpose-convolution overlap') | |
| previous = state.index_select(0, slots) | |
| tail = y[..., produced:] | |
| if module.conv.bias is not None: | |
| tail = tail - module.conv.bias[None, :, None] | |
| state.index_copy_(0, slots, tail) | |
| y[..., :overlap] += previous # disjoint from the tail because overlap <= produced | |
| return y[..., :produced] | |
| if isinstance(module, MimiResnetBlock): | |
| shortcut = self._layer(module.shortcut, x, slots, key + '.shortcut') | |
| for i, child in enumerate(module.block): | |
| x = self._layer(child, x, slots, key + f'.block.{i}') | |
| return shortcut + x | |
| return module(x) | |
| def _rvq(rvq, codes): | |
| total = None | |
| for i in range(codes.shape[1]): | |
| e = F.embedding(codes[:, i], rvq.layers[i].codebook.embed).permute(0, 2, 1) | |
| total = e if total is None else total + e | |
| return rvq.output_proj(total) if rvq.output_proj is not None else total | |
| def _quantizer(self, codes): | |
| # Same arithmetic as MimiSplitResidualVectorQuantizer.decode, without its in-graph torch.tensor(0.0). | |
| q = self.m.quantizer | |
| n = q.num_semantic_quantizers | |
| out = self._rvq(q.semantic_residual_vector_quantizer, codes[:, :n]) | |
| if codes.shape[1] > n: | |
| out = out + self._rvq(q.acoustic_residual_vector_quantizer, codes[:, n:]) | |
| return out | |
| def _transformer(self, h, slots): | |
| batch, count, _ = h.shape | |
| start = self.tokens.index_select(0, slots) | |
| positions = start[:, None] + self._steps[None] # [B,T] absolute | |
| ring_index = positions % self.ring | |
| rows = slots[:, None] | |
| self.key_pos.index_put_((rows, ring_index), positions) | |
| key_pos = self.key_pos.index_select(0, slots)[:, None, :] # [B,1,R] | |
| query_pos = positions[:, :, None] # [B,T,1] | |
| mask = ((key_pos <= query_pos) & (key_pos > query_pos - self.window))[:, None] # [B,1,T,R] | |
| freqs = positions[:, :, None].float() * self._inv_freq[None, None, :] | |
| emb = torch.cat((freqs, freqs), dim=-1) | |
| cos = (emb.cos() * self._rope_scale).to(h.dtype) | |
| sin = (emb.sin() * self._rope_scale).to(h.dtype) | |
| for li, layer in enumerate(self.m.decoder_transformer.layers): | |
| attn = layer.self_attn | |
| x = layer.input_layernorm(h) | |
| q = attn.q_proj(x).view(batch, count, attn.num_heads, self.head_dim).transpose(1, 2) | |
| k = attn.k_proj(x).view(batch, count, self.kv_heads, self.head_dim).transpose(1, 2) | |
| v = attn.v_proj(x).view(batch, count, self.kv_heads, self.head_dim).transpose(1, 2) | |
| q, k = apply_rotary_pos_emb(q, k, cos, sin) | |
| keys, values = self.kv[li, 0], self.kv[li, 1] | |
| # Advanced indices on dims 0 and 2 put the indexed dims first: value layout [B,T,heads,dim]. | |
| keys[rows, :, ring_index] = k.transpose(1, 2) | |
| values[rows, :, ring_index] = v.transpose(1, 2) | |
| k_all = repeat_kv(keys.index_select(0, slots), self.kv_groups) | |
| v_all = repeat_kv(values.index_select(0, slots), self.kv_groups) | |
| a = F.scaled_dot_product_attention(q, k_all, v_all, attn_mask=mask) | |
| a = attn.o_proj(a.transpose(1, 2).reshape(batch, count, -1)) | |
| h = h + layer.self_attn_layer_scale(a) | |
| h = h + layer.mlp_layer_scale(layer.mlp(layer.post_attention_layernorm(h))) | |
| self.tokens.index_copy_(0, slots, start + count) | |
| self.tokens[self.dummy].zero_() | |
| return h | |
| def _forward(self, slots, codes): | |
| x = self._quantizer(codes) | |
| x = self._layer(self.m.upsample, x, slots, 'upsample') | |
| x = self._transformer(x.transpose(1, 2), slots).transpose(1, 2) | |
| for i, layer in enumerate(self.m.decoder.layers): | |
| x = self._layer(layer, x, slots, f'decoder.{i}') | |
| return x | |
| # ---- capture ------------------------------------------------------------------------------ | |
| def _capture(self, size): | |
| bucket = _Bucket() | |
| bucket.slots = torch.full((size,), self.dummy, device=self.device, dtype=torch.long) | |
| bucket.codes = torch.zeros(size, self.codebooks, 1, device=self.device, dtype=torch.long) | |
| bucket.pinned = [torch.full((size,), self.dummy, dtype=torch.long).pin_memory() for _ in range(8)] | |
| bucket.events = [None] * len(bucket.pinned) | |
| bucket.turn = 0 | |
| # Warmup and capture touch only the dummy slot, so live seats are never disturbed. | |
| side = torch.cuda.Stream(device=self.device) | |
| side.wait_stream(torch.cuda.current_stream(self.device)) | |
| with torch.cuda.stream(side): | |
| for _ in range(3): | |
| self._forward(bucket.slots, bucket.codes) | |
| torch.cuda.current_stream(self.device).wait_stream(side) | |
| torch.cuda.synchronize(self.device) | |
| if self._pool is None: | |
| self._pool = torch.cuda.graph_pool_handle() | |
| bucket.graph = torch.cuda.CUDAGraph() | |
| with torch.cuda.graph(bucket.graph, pool=self._pool): | |
| bucket.out = self._forward(bucket.slots, bucket.codes) | |
| if self.samples_per_frame is None: | |
| self.samples_per_frame = bucket.out.shape[-1] | |
| self._reset_row(self.dummy) | |
| self._buckets[size] = bucket | |
| return bucket | |
| def warmup(self, batch_buckets=None): | |
| torch.backends.cudnn.benchmark = False | |
| sizes = self.batch_buckets if batch_buckets is None else tuple(sorted(set(int(b) for b in batch_buckets))) | |
| for size in sorted(sizes, reverse=True): # largest first so the shared pool is sized once | |
| if size not in self._buckets: | |
| self._capture(size) | |
| self.batch_buckets = tuple(sorted(set(self.batch_buckets) | set(sizes))) | |
| # ---- public ------------------------------------------------------------------------------- | |
| def _reset_row(self, row): | |
| self.conv_flat[row].zero_() | |
| self.key_pos[row].fill_(NEVER) | |
| self.tokens[row].zero_() | |
| def reset(self, slot, zero_kv=False): | |
| """Three fill kernels, no host sync. KV rows are unreachable once key_pos is NEVER | |
| (masked weights are exactly zero); zero_kv=True additionally clears them.""" | |
| if not 0 <= slot < self.max_slots: | |
| raise IndexError(f'slot {slot} outside [0,{self.max_slots})') | |
| self._reset_row(slot) | |
| if zero_kv: | |
| self.kv[:, :, slot].zero_() | |
| def decode(self, slots, codes): | |
| """slots: distinct ints in [0,max_slots) (or a CUDA long tensor of them); codes: [B,K] long CUDA, | |
| one frame per seat. Returns an owned float32 [B, samples_per_frame] tensor.""" | |
| count = len(slots) | |
| if codes.shape[0] != count or codes.dim() != 2 or codes.shape[1] != self.codebooks: | |
| raise ValueError(f'codes must be [{count},{self.codebooks}], got {tuple(codes.shape)}') | |
| size = next((b for b in self.batch_buckets if b >= count), None) | |
| if size is None: | |
| raise ValueError(f'batch {count} exceeds largest bucket {self.batch_buckets[-1]}') | |
| bucket = self._buckets.get(size) or self._capture(size) | |
| if torch.is_tensor(slots): | |
| bucket.slots[:count].copy_(slots) | |
| bucket.slots[count:].fill_(self.dummy) | |
| else: | |
| if len(set(slots)) != count or min(slots) < 0 or max(slots) >= self.max_slots: | |
| raise ValueError('slots must be distinct and inside [0,max_slots)') | |
| # Rotate pinned staging rows so an async H2D copy is never overwritten before it runs. | |
| turn = bucket.turn | |
| bucket.turn = (turn + 1) % len(bucket.pinned) | |
| if bucket.events[turn] is not None: | |
| bucket.events[turn].synchronize() | |
| staging = bucket.pinned[turn] | |
| staging[:count] = torch.tensor(slots, dtype=torch.long) | |
| staging[count:] = self.dummy | |
| bucket.slots.copy_(staging, non_blocking=True) | |
| event = torch.cuda.Event() | |
| event.record() | |
| bucket.events[turn] = event | |
| bucket.codes[:count, :, 0].copy_(codes) | |
| bucket.graph.replay() | |
| return bucket.out[:count, 0].float().clone() if bucket.out.dtype != torch.float32 else bucket.out[:count, 0].clone() | |