GB10 real-time serving: opt-in quantization + serve-path optimizations (2.2x: 101.7 -> 46.5 ms/frame)

moshi/moshi/fp8_quantize.py

@@ -0,0 +1,244 @@
+"""FP8 dynamic quantization for PersonaPlex/Moshi.
+
+Usage:
+    from moshi.fp8_quantize import quantize_model, free_bf16_inproj
+
+    lm = loaders.get_moshi_lm(...)
+    quantize_model(lm)       # Quantize weights + patch forward paths
+    # ... warmup ...
+    free_bf16_inproj(lm)     # Free original bf16 in_proj copies
+
+Replaces nn.Linear weights with FP8 (float8_e4m3fn) and patches the
+gating/attention forward paths to use torch._scaled_mm for native FP8 GEMM.
+
+Benchmarked results:
+    - Jetson Thor:  114ms → 74ms lm_step (1.54x), 16.7 → 11.2 GB
+    - DGX Spark:    95.7ms → 70ms lm_step (1.37x), 18.8 → 11.0 GB
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import logging
+import types
+
+logger = logging.getLogger(__name__)
+
+
+def fp8_linear(x, weight_fp8, weight_scale, bias=None):
+    """FP8 GEMM with dynamic per-tensor input scaling."""
+    orig_shape = x.shape
+    in_features = weight_fp8.shape[1]
+    out_features = weight_fp8.shape[0]
+    x_2d = x.reshape(-1, in_features)
+
+    x_amax = x_2d.abs().amax()
+    x_scale = (x_amax / 448.0).clamp(min=1e-12)
+    x_fp8 = (x_2d / x_scale).to(torch.float8_e4m3fn)
+
+    out = torch._scaled_mm(
+        x_fp8, weight_fp8.t(),
+        scale_a=x_scale.float().view(1),
+        scale_b=weight_scale,
+        out_dtype=torch.bfloat16
+    )
+
+    if bias is not None:
+        out = out + bias
+    return out.reshape(*orig_shape[:-1], out_features)
+
+
+# ============================================================================
+# Module-level forward patch for nn.Linear
+# ============================================================================
+
+def _fp8_forward(self, x):
+    """Patched nn.Linear.forward using dynamic FP8."""
+    return fp8_linear(x, self.weight, self.weight_scale, self.bias)
+
+
+# ============================================================================
+# Gating forward patch (ActivationGating uses F.linear, not nn.Linear.forward)
+# ============================================================================
+
+def _make_gating_forward():
+    def gating_forward_fp8(self, x):
+        lin_in = self.linear_in
+        lin_out = self.linear_out
+        if getattr(lin_in, '_is_fp8', False):
+            x = fp8_linear(x, lin_in.weight, lin_in.weight_scale)
+        else:
+            x = F.linear(x, lin_in.weight)
+        B, T, _ = x.shape
+        x = x.view(B, T, 2, -1)
+        x = self.activation(x[..., 0, :]) * x[..., 1, :]
+        if getattr(lin_out, '_is_fp8', False):
+            x = fp8_linear(x, lin_out.weight, lin_out.weight_scale)
+        else:
+            x = F.linear(x, lin_out.weight)
+        return x
+    return gating_forward_fp8
+
+
+# ============================================================================
+# Attention forward patch (in_proj_weight is bare nn.Parameter, not nn.Linear)
+# ============================================================================
+
+def _make_attn_forward():
+    from einops import rearrange as _rearrange
+
+    def attn_forward_fp8(self, query, key, value):
+        import moshi.modules.transformer as tf_mod
+
+        state = self._streaming_state
+        T = query.shape[1]
+
+        if state is None:
+            offset = torch.zeros(1, device=query.device, dtype=torch.long)
+            offset_cpu = 0
+        else:
+            offset = state.offset
+            offset_cpu = state.offset_cpu
+
+        if self.weights_per_step:
+            projected = tf_mod.multi_linear(
+                self.weights_per_step, self.in_proj_weight, query, offset_cpu
+            )
+        else:
+            if getattr(self, '_in_proj_fp8', False):
+                projected = fp8_linear(query, self._in_proj_fp8_weight, self._in_proj_scale)
+            else:
+                projected = F.linear(query, self.in_proj_weight)
+
+        q, k, v = _rearrange(
+            projected, "b t (p h d) -> p b h t d", p=3, h=self.num_heads
+        )
+
+        if self.rope:
+            q, k = self.rope(q, k, offset, time_before_heads=False)
+
+        k, v, pos_k = self._complete_kv(k, v)
+        if self.causal:
+            pos_k = pos_k.view(1, -1)
+            pos_q = offset + torch.arange(
+                T, device=query.device, dtype=torch.long
+            ).view(-1, 1)
+            delta = pos_q - pos_k
+            attn_bias = (pos_k >= 0) & (delta >= 0)
+            if self.context is not None:
+                attn_bias = attn_bias & (delta < self.context)
+        else:
+            attn_bias = None
+        x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
+
+        x = _rearrange(x, "b h t d -> b t (h d)")
+        if self.weights_per_step:
+            x = tf_mod.multi_linear(
+                self.weights_per_step, self.out_proj.weight, x, offset_cpu
+            )
+        else:
+            out_proj = self.out_proj
+            if getattr(out_proj, '_is_fp8', False):
+                x = fp8_linear(x, out_proj.weight, out_proj.weight_scale)
+            else:
+                x = out_proj(x)
+        if state is not None:
+            state.offset.add_(T)
+            state.offset_cpu += T
+        return x
+
+    return attn_forward_fp8
+
+
+# ============================================================================
+# Weight quantization
+# ============================================================================
+
+def quantize_linear_fp8(module):
+    """Quantize a single nn.Linear to FP8 in-place."""
+    w = module.weight.data
+    amax = w.abs().amax()
+    scale = (amax / 448.0).clamp(min=1e-12)
+    module.weight = nn.Parameter(
+        (w / scale).to(torch.float8_e4m3fn), requires_grad=False
+    )
+    module.register_buffer('weight_scale', scale.float().view(1))
+    module._is_fp8 = True
+    module.forward = types.MethodType(_fp8_forward, module)
+
+
+def quantize_model(model, min_features=512):
+    """Quantize all large Linear layers in the model to FP8.
+
+    Patches ActivationGating and StreamingMultiheadAttention forward methods
+    to use FP8 GEMM via torch._scaled_mm.
+
+    Skips depformer self_attn (matrices too small — overhead > savings).
+    """
+    import moshi.modules.gating as gating_mod
+    import moshi.modules.transformer as tf_mod
+
+    gating_mod.ActivationGating.forward = _make_gating_forward()
+    tf_mod.StreamingMultiheadAttention.forward = _make_attn_forward()
+
+    linear_count = 0
+    for name, module in list(model.named_modules()):
+        if isinstance(module, nn.Linear):
+            if module.in_features < min_features and module.out_features < min_features:
+                continue
+            if module.weight.ndim > 2:
+                continue
+            # Skip depformer self_attn — per-step slices too small for FP8 benefit
+            if 'depformer' in name and 'self_attn' in name:
+                continue
+            quantize_linear_fp8(module)
+            linear_count += 1
+
+    # Quantize bare in_proj_weight parameters on main transformer attention
+    inproj_count = 0
+    for name, module in list(model.named_modules()):
+        if isinstance(module, tf_mod.StreamingMultiheadAttention):
+            if 'depformer' in name:
+                continue
+            if module.weights_per_step:
+                continue
+            w = module.in_proj_weight.data
+            if w.ndim != 2:
+                continue
+            amax = w.abs().amax()
+            scale = (amax / 448.0).clamp(min=1e-12)
+            module.register_buffer('_in_proj_fp8_weight',
+                (w / scale).to(torch.float8_e4m3fn))
+            module._in_proj_scale = scale.float().view(1)
+            module._in_proj_fp8 = True
+            inproj_count += 1
+
+    logger.info(f"[FP8] Quantized {linear_count} Linear + {inproj_count} in_proj_weight")
+    logger.info(f"[FP8] GPU memory: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
+    return model
+
+
+# ============================================================================
+# Cleanup — free original bf16 copies after CUDA graph warmup
+# ============================================================================
+
+def free_bf16_inproj(model):
+    """Free bf16 in_proj_weight copies after warmup (KV cache already allocated).
+
+    Call this after state.warmup() — the CUDA graph has captured the FP8 path,
+    so the original bf16 weights are no longer needed.
+    """
+    import moshi.modules.transformer as tf_mod
+    freed = 0
+    for name, module in model.named_modules():
+        if isinstance(module, tf_mod.StreamingMultiheadAttention):
+            if getattr(module, '_in_proj_fp8', False):
+                sz = module.in_proj_weight.numel() * module.in_proj_weight.element_size()
+                module.in_proj_weight = nn.Parameter(
+                    torch.empty(0, dtype=torch.bfloat16, device=module.in_proj_weight.device),
+                    requires_grad=False
+                )
+                freed += sz
+    if freed > 0:
+        torch.cuda.empty_cache()
+        logger.info(f"[FP8] Freed {freed/1e9:.2f} GB bf16 in_proj_weight copies")
+        logger.info(f"[FP8] GPU memory after cleanup: {torch.cuda.memory_allocated() / 1e9:.2f} GB")

moshi/moshi/models/lm.py

@@ -176,6 +176,8 @@ def encode_from_sphn(mimi, samples, max_batch=sys.maxsize):
             break
 
         batch = torch.cat(current_batch, dim=0)  # shape: (B, C, T)
+        model_dtype = next(mimi.parameters()).dtype
+        batch = batch.to(dtype=model_dtype)
         encoded = mimi.encode(batch)  # shape: (B, K, F)
         separated = torch.unbind(encoded, dim=0)  # shape: (K, F)
         reshaped = [x.unsqueeze(0) for x in separated]  # shape: (1, K, F)
@@ -661,6 +663,7 @@ def __init__(
         save_voice_prompt_embeddings: bool = False,
         sample_rate: int = 32000,
         frame_rate: int = FRAME_RATE_HZ,
+        depformer_early_exit: Optional[int] = None,
     ):
         assert not lm_model.training, "generation shouldn't be used in training mode."
         super().__init__()
@@ -693,6 +696,20 @@ def __init__(
         self.delays_cuda = torch.tensor(
             lm_model.delays, device=lm_model.device, dtype=torch.long
         )
+        self.depformer_early_exit = depformer_early_exit
+        if depformer_early_exit is not None:
+            # The serve stack decodes agent audio from codebooks 1..8, so at
+            # least 8 depformer steps must run. Codebooks beyond the exit
+            # point are only skippable because every step of the serve flow
+            # (warmup, voice/text prompts, streaming) provides the user-side
+            # tokens, which overwrite the sampled values in the cache.
+            assert 8 <= depformer_early_exit <= lm_model.dep_q, (
+                f"depformer_early_exit must be in [8, {lm_model.dep_q}]"
+            )
+            assert not return_logits and not report_loss, (
+                "depformer_early_exit is incompatible with "
+                "return_logits/report_loss"
+            )
         self.save_voice_prompt_embeddings = save_voice_prompt_embeddings
         self.voice_prompt_audio: Optional[torch.Tensor] = None
         self.voice_prompt_cache: Optional[torch.Tensor] = None
@@ -817,6 +834,12 @@ def step(self, input_tokens: torch.Tensor=None, moshi_tokens:torch.Tensor=None,
         -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, dict[str, torch.Tensor]]:
         state = self._streaming_state
         lm_model = self.lm_model
+        if self.depformer_early_exit is not None and input_tokens is None:
+            raise RuntimeError(
+                "depformer_early_exit requires input_tokens (user audio codes) "
+                "to be provided on every step: the skipped codebooks are only "
+                "safe to omit when provided tokens overwrite them."
+            )
         prepared_inputs = self.prepare_step_input(
             input_tokens, moshi_tokens, text_token,
         )
@@ -1139,8 +1162,11 @@ def depformer_step(
         depformer_tokens: list[torch.Tensor] = []
         depformer_logits: list[torch.Tensor] = []
         assert not lm_model.depformer.is_streaming
+        n_steps = lm_model.dep_q
+        if self.depformer_early_exit is not None:
+            n_steps = min(self.depformer_early_exit, lm_model.dep_q)
         with lm_model.depformer.streaming(B):
-            for cb_index in range(lm_model.dep_q):
+            for cb_index in range(n_steps):
                 input_ = prev_token[:, None, None]
                 logits = lm_model.forward_depformer(cb_index, input_, transformer_out)
                 if self.return_logits:
@@ -1163,6 +1189,14 @@ def depformer_step(
                 )
                 depformer_tokens.append(next_token)
 
+        if len(depformer_tokens) < lm_model.dep_q:
+            # Early exit: pad the skipped codebooks with zero_token_id (-1,
+            # "no input"). These values never reach the cache in the serve
+            # flow because the corresponding channels are always provided.
+            pad = torch.full_like(depformer_tokens[0], lm_model.zero_token_id)
+            depformer_tokens.extend(
+                [pad] * (lm_model.dep_q - len(depformer_tokens))
+            )
         assert len(depformer_tokens) == lm_model.dep_q, (
             len(depformer_tokens),
             lm_model.dep_q,

moshi/moshi/models/loaders.py

@@ -252,6 +252,19 @@ def get_moshi_lm(
 
     # Assign weights to target device
     dev = torch.device(device) if isinstance(device, str) else device
+
+    # Patch 3: materialize any model parameters still absent from the
+    # checkpoint. The model is initialized on the meta device, so a key that
+    # is neither in the checkpoint nor covered by the patches above would
+    # survive load_state_dict(strict=False, assign=True) as a meta tensor and
+    # crash the final .to() with "Cannot copy out of meta tensor". This
+    # happens e.g. for depformer_emb.7 when expanding a base Moshi dep_q=8
+    # checkpoint to dep_q=16 (indices 8..15 are backfilled from 0..7, but
+    # index 7 itself does not exist in the checkpoint). Zero-init and warn.
+    for name, tensor in model_sd.items():
+        if name not in state_dict:
+            print(f"Zero-initializing key missing from checkpoint: {name}")
+            state_dict[name] = torch.zeros(tensor.shape, dtype=dtype, device=dev)
     for key in state_dict:
         state_dict[key] = state_dict[key].to(device=dev, dtype=dtype)