Fix H200 fp16 instability: pin MATH SDPA, fp32 autoencoder, no TF32

The continuation sounds clean on local Blackwell (sm_120) but produced
sporadic loud bursts on ZeroGPU's H200 (sm_90). Root cause is fp16
numerical instability on H200, not a torch version.

- Pin F.scaled_dot_product_attention to the MATH (reference) backend
around generation. SA3's transformer calls SDPA with no backend pin and
casts q/k/v to fp16; on H200 torch 2.8+ can route that to a cuDNN/
mem-efficient kernel with known fp16 NaN/garbage bugs (the loud-burst
signature). flash-attn is absent in this image, so SDPA is the live
path and the pin engages.
- Run the autoencoder (Oobleck/DAC, Snake1d) in fp32 while the DiT stays
fp16. Snake's x + (1/(b+1e-9))*sin(a*x)^2 can exceed fp16's 65504
ceiling -> inf/NaN -> broadband noise on decode. sampling.py casts
latents to the pretransform dtype before decode, so an fp32 pretransform
upcasts the whole encode+decode path with no library change. The inpaint
lead is cast to the pretransform dtype to match.
- Disable TF32 matmul/cudnn accumulation, which compounds error through
the 8-step distilled sampler.

Verified end-to-end on Blackwell: DiT fp16 / pretransform fp32, output
finite, peak 1.0, clean burst/flatness/hf metrics, no regression.

engine.py

@@ -37,11 +37,43 @@ Mask convention (verified against the installed library source):
   at the front and mask [lead, lead+new], so SA3 keeps the lead and generates a
   fresh `new`-second tail that continues from the clip's end.
 """
+import contextlib
 import time
 
 import numpy as np
 import torch
 
+# --- H200 (sm_90) fp16 numerical-stability hardening -------------------------
+# CODA sounds perfect on local Blackwell (sm_120) but produced "sporadic loud
+# bursts" on ZeroGPU's H200. Root cause is fp16 numerical instability on H200,
+# not a torch version: (1) torch 2.8 changed cuDNN SDPA backend selection for
+# H200 and the cuDNN/mem-efficient fp16 SDPA kernels have known NaN/garbage bugs
+# (pytorch#139298/#166211/#124877/#112577) — exactly the "loud burst" signature;
+# (2) TF32 matmul accumulation adds error that compounds through the 8-step
+# distilled sampler. We pin the math (reference) SDPA backend around generation
+# and force full-precision matmul accumulation. Both are no-ops on CPU.
+torch.backends.cuda.matmul.allow_tf32 = False
+torch.backends.cudnn.allow_tf32 = False
+
+try:
+    from torch.nn.attention import SDPBackend, sdpa_kernel as _sdpa_kernel
+except Exception:                                    # very old/unexpected torch
+    SDPBackend = None
+    _sdpa_kernel = None
+
+
+def _sdpa_math():
+    """Context manager that forces F.scaled_dot_product_attention onto the MATH
+    (reference) backend for the wrapped call. SA3's transformer calls SDPA with
+    no backend pin and casts q/k/v to fp16; on H200 torch may route that to a
+    cuDNN/mem-efficient kernel with the known fp16 NaN bug. MATH is the stable
+    reference path. Returns a fresh context each call; a no-op if the API is
+    missing (older torch) or on CPU."""
+    if _sdpa_kernel is None:
+        return contextlib.nullcontext()
+    return _sdpa_kernel(SDPBackend.MATH)
+
+
 MODEL_ID = "stabilityai/stable-audio-3-small-music"
 
 SR = 44100               # SA3 native sample rate (model_config: sample_rate)
@@ -102,16 +134,28 @@ def preload():
 
 
 def _ensure_on_device():
-    """Ensure weights are on the GPU in fp16. Called inside the @spaces.GPU
-    window on every generation. `.to()` is a cheap no-op when the model is
-    already placed, so we re-ensure each call rather than caching device state —
-    that stays correct even if ZeroGPU detaches the GPU between calls. fp16 is
-    only valid on CUDA; on CPU the model stays fp32."""
+    """Ensure weights are on the GPU with the right per-component dtype. Called
+    inside the @spaces.GPU window on every generation. `.to()` is a cheap no-op
+    when already placed, so we re-ensure each call rather than caching device
+    state — that stays correct even if ZeroGPU detaches the GPU between calls.
+
+    Dtype split (the H200 quality fix): the DiT runs in fp16 (fast, and Stability
+    loads the model fp16), but the autoencoder/pretransform runs in fp32. SA3's
+    Oobleck/DAC decoder uses Snake1d activations — x + (1/(beta+1e-9))*sin(a*x)^2
+    — whose reciprocal*sin^2 term can exceed fp16's ~65504 ceiling, giving
+    inf/NaN -> loud broadband noise on decode. fp32 there has ample headroom.
+    sampling.py casts the sampled latents to the pretransform's dtype right
+    before `decode`, so an fp32 pretransform makes the WHOLE encode+decode path
+    fp32 with no library change; only the DiT stays fp16. All fp16/fp32 here is
+    CUDA-only; on CPU the model stays fp32 throughout."""
     global _model, _on_device
     dev = _device()
     _model = _model.to(dev)
     if dev == "cuda":
         _model = _model.to(torch.float16)
+        # keep the autoencoder in fp32 (Snake1d fp16 overflow guard)
+        if getattr(_model, "pretransform", None) is not None:
+            _model.pretransform.to(torch.float32)
     _on_device = dev
     return _model
 
@@ -260,10 +304,14 @@ def continue_audio(clip_path, total_seconds, prompt="", cfg_scale=DEFAULT_CFG,
     # a long source no longer fills the buffer and starves the masked region.
     lead_samples = min(int(round(lead * SR)), source.shape[-1])
     lead_audio = source[:, -lead_samples:]
-    # the autoencoder runs in the model's dtype (fp16 on CUDA); the conditioning
-    # audio must match it or the encoder's conv1d rejects the input dtype.
-    model_dtype = next(model.model.parameters()).dtype
-    lead_audio = lead_audio.to(model_dtype)
+    # the lead is encoded by the autoencoder (pretransform), which we now run in
+    # fp32; its conv1d rejects a mismatched input dtype, so match the
+    # pretransform's parameter dtype (fp32 on CUDA), NOT the DiT's fp16. Resample
+    # is a no-op here (clip is already 44.1k), so prepare_audio keeps this dtype.
+    ae = getattr(model, "pretransform", None)
+    ae_dtype = (next(ae.parameters()).dtype if ae is not None
+                else next(model.model.parameters()).dtype)
+    lead_audio = lead_audio.to(ae_dtype)
 
     mask_start, mask_end = lead, buffer_seconds
     prompt = (prompt or "").strip()
@@ -275,7 +323,9 @@ def continue_audio(clip_path, total_seconds, prompt="", cfg_scale=DEFAULT_CFG,
 
     def _draw(draw_seed):
         """One full SA3 inpaint draw -> normalized generated tail (2, M)."""
-        with torch.no_grad():
+        # _sdpa_math() forces the reference attention backend for the whole
+        # sample loop — the H200 fp16-SDPA-garbage guard (no-op on CPU).
+        with torch.no_grad(), _sdpa_math():
             output = generate_diffusion_cond_inpaint(
                 model,
                 steps=STEPS,