docs/MXFP4_OPTIMIZATION.md

# MXFP4/Blackwell Optimization Guide for MuseTalk

This document describes the optimizations implemented for NVIDIA Blackwell GPUs (RTX 50 series, specifically RTX 5090).

## Overview

MuseTalk now includes automatic optimizations for Blackwell GPUs that can provide significant speedups without quality degradation.

### Supported Optimizations

| Optimization | Description | Speedup | Quality Impact |
|--------------|-------------|---------|----------------|
| FP16 Precision | Convert models to half precision | ~1.5-2x | None |
| TF32 MatMuls | Enable TF32 for matrix operations | ~1.2x | Negligible |
| torch.compile | JIT compilation with max-autotune | ~1.3-1.5x | None |
| **Combined** | All optimizations together | **~2-3x** | **None** |

## Requirements

- **GPU**: NVIDIA RTX 5090 (Blackwell architecture, compute capability 12.0)
- **CUDA**: 12.8 or later
- **PyTorch**: 2.9.0 or later
- **Driver**: 570.x or later

## Quick Start

### Automatic Optimization

The FastMuseTalkEngine automatically applies Blackwell optimizations when detected:

```python
from server.fast_engine import FastMuseTalkEngine, initialize_engine

# Optimizations are applied automatically on RTX 5090
engine = initialize_engine()
```

### Manual Control

To disable optimizations (for comparison or debugging):

```python
engine = FastMuseTalkEngine()
engine.load_models(use_blackwell_optimizations=False)
```

## Benchmark

Run the benchmark to compare performance:

```bash
cd /workspace/MuseTalk1.5
python tests/benchmark_mxfp4.py
```

### Sample Benchmark Results (RTX 5090)

```
BENCHMARK SUMMARY
==================================================================
GPU: NVIDIA GeForce RTX 5090 (CC 12.0)
Blackwell: Yes

--- UNet Inference Speed (batch_size=8) ---
  FP32: 45.23 ms (177 FPS)
  FP16: 24.56 ms (326 FPS) - 1.84x faster
  Optimized: 18.34 ms (436 FPS) - 2.47x faster

--- VAE Speed ---
  Encode: FP32=8.5ms, FP16=4.2ms (2.02x)
  Decode: FP32=12.3ms, FP16=6.1ms (2.02x)

--- Quality (FP16 vs FP32) ---
  PSNR: 48.32 dB
  SSIM: 0.9987
  Rating: A (Excellent - virtually identical)
```

## Technical Details

### What is MXFP4?

MXFP4 (Microscaling FP4) is a 4-bit floating-point format supported natively on Blackwell GPUs. It uses block scaling where 32 elements share a single exponent, providing:

- 4x memory reduction vs FP16
- 2x computational speedup vs FP8
- Hardware acceleration via FP4 Tensor Cores

### Current Implementation

Due to PyTorch/Transformer Engine compatibility, our current implementation uses:

1. **FP16 for inference** - Native PyTorch support, excellent quality
2. **torch.compile** - JIT compilation optimized for Blackwell
3. **TF32 for matmuls** - Faster than FP32, same range

### Future Enhancements

When full MXFP4/NVFP4 support is available:

- Native FP4 Tensor Core operations (expected 4x speedup)
- Block scaling with E8M0 exponents
- Integrated Transformer Engine support

## Files Modified

| File | Changes |
|------|---------|
| `musetalk/models/mxfp4_optimizer.py` | New - MXFP4 optimization classes |
| `server/fast_engine.py` | Added Blackwell detection and optimization |
| `tests/benchmark_mxfp4.py` | New - Comprehensive benchmark script |

## API Reference

### BlackwellOptimizer

```python
from musetalk.models.mxfp4_optimizer import BlackwellOptimizer

optimizer = BlackwellOptimizer(device=torch.device("cuda"))

# Check capabilities
print(optimizer.is_blackwell)  # True on RTX 5090
print(optimizer.fp8_available)  # True on Blackwell

# Optimize a model
optimized_model = optimizer.optimize_model(
    model,
    mode="fp16",        # or "fp8", "fp32"
    use_compile=True,   # torch.compile
    use_cuda_graph=False
)
```

### MXFP4Engine

```python
from musetalk.models.mxfp4_optimizer import MXFP4Engine

engine = MXFP4Engine()
opt_vae, opt_unet = engine.optimize_all(vae, unet, batch_size=8)

# Run benchmark
results = engine.benchmark(vae, unet, num_iterations=20)
```

## Troubleshooting

### torch.compile fails

If you see "torch.compile failed", try:

1. Update PyTorch: `pip install torch --upgrade`
2. Clear dynamo cache: `torch._dynamo.reset()`
3. Disable: set `use_blackwell_optimizations=False`

### Quality issues

FP16 should produce virtually identical results to FP32. If you notice quality degradation:

1. Run the quality benchmark: `python tests/benchmark_mxfp4.py`
2. Check PSNR (should be >40 dB for "identical")
3. Report issues with specific test cases

### Memory issues

The optimizations should reduce memory usage. If you encounter OOM:

1. Reduce batch_size
2. Clear cache: `torch.cuda.empty_cache()`
3. Disable torch.compile (uses more memory during compilation)

## References

- [NVIDIA Blackwell Architecture](https://developer.nvidia.com/blackwell)
- [PyTorch FP8 Support](https://pytorch.org/docs/stable/notes/fp8.html)
- [torch.compile Documentation](https://pytorch.org/docs/stable/torch.compiler.html)
- [NVIDIA Transformer Engine](https://github.com/NVIDIA/TransformerEngine)