Ampere (sm_86) compatibility for fp8-cast LoRA fusion and mono-audio input in a2vid

[Larry Bowden (fastflair)]

Problem

Three issues prevent the distilled + a2vid pipelines from running on Ampere
GPUs (sm_86, e.g. RTX 3090) with --offload cpu and mono audio input:

1. _read_scales mmaps the full checkpoint to enumerate key names
safetensors.safe_open() on a 46 GB bf16 checkpoint requires the OS to
commit a 46 GB virtual-memory region. On hosts where total committable
memory (RAM + swap) is below that — common on WSL2 and consumer Linux
machines — this raises Cannot allocate memory (12) during argument
parsing, before any model is loaded. A bf16 checkpoint has no *_scale
tensors, so the entire mmap was spent on h.keys() returning an empty
result.

2. fuse_cast_fp8_weight uses a Triton kernel that doesn't compile on Ampere
When the distilled LoRA is fused into fp8-cast weights (required for
dialogue/a2vid shots), fused_add_round_kernel casts to fp8e4nv
(float8_e4m3fn). Triton only supports this dtype on Ada (sm_89) and
newer; on Ampere it raises:
CompilationError: type fp8e4nv not supported in this architecture.
The supported fp8 dtypes are ('fp8e4b15', 'fp8e5')
A deterministic bf16 fallback already exists in fuse_cast_fp8_weight
but was unreachable because the gate only checked TRITON_AVAILABLE, not
GPU capability.

3. decode_audio_from_file returns mono for mono input; a2vid requires stereo
The audio VAE's conv_in has in_channels=2 and
_validate_audio_waveform rejects mono. When TTS or any mono wav is
passed via --audio-path, the pipeline crashes with:
RuntimeError: expected input[1, 1, 274, 66] to have 2 channels, but got 1

---

Fix

fp8_cast.py

• _read_scales: read only the safetensors header (8-byte length prefix +
  JSON) to enumerate key names, avoiding the full-file mmap. The
  safe_open path is still taken when *_scale keys are actually present
  (prequant fp8 checkpoints). Behavior is identical for all checkpoint
  types; only the wasted mapping on bf16 is removed.
• Add _cuda_cc() (cached compute capability) and
  _triton_can_stochastic_round(dtype): returns True only when Triton
  can compile the fp8 kernel for the given dtype on the current GPU
  (float8_e4m3fn requires sm_89+; float8_e5m2 requires sm_80+). Gate
  the fused_add_round_launch call on this instead of bare
  TRITON_AVAILABLE. On unsupported hardware, the existing deterministic
  bf16 fallback is taken — no quality change for inference, no rounding
  difference for prequant-fp8 checkpoints.

media_io.py

• decode_audio_from_file: upmix mono waveforms to dual-mono stereo
  (np.repeat(audio, 2, axis=0)) before returning. No-op for already-
  stereo input.

---

Testing

• RTX 3090 (sm_86), WSL2, 31 GiB RAM + 71 GiB swap, --offload cpu,
  --quantization fp8-cast, distilled checkpoint (bf16), distilled LoRA
• Silent shots (distilled pipeline): _read_scales fix eliminates the
  mmap crash; shots render end-to-end
• Dialogue shots (a2vid_two_stage + LoRA fusion): all three fixes required;
  shots render end-to-end after patch
• Ada / no regression expected: _triton_can_stochastic_round returns
  True on sm_89+, preserving stochastic rounding on capable hardware
  (untested by author — please verify on your CI)

---

Notes

The _read_scales change does not affect prequant fp8 checkpoints: the
early-exit if not scale_keys: return out is not taken, and the full
safe_open path runs as before. The header read is O(KB) regardless of
checkpoint size, so there is no measurable overhead on the hot path.

The stereo upmix produces L=R dual-mono, which is appropriate for speech
input. The audio VAE and mux path treat both channels identically, so
there is no stereo-width artifact.