serve/fast_codec.py

"""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)

    @staticmethod
    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

    @torch.inference_mode()
    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_()

    @torch.inference_mode()
    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_()

    @torch.inference_mode()
    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()