iGavroche
20626b4fd4
- ROCMSamplerCustomAdvanced: V3 ComfyNode drop-in for stock SamplerCustomAdvanced with architecture detection, ROCm backend tuning, memory defrag for high-memory models (LTX 128ch), enhanced per-step callback with ETA - ROCMSamplerCustomAdvancedBenchmark: A/B comparison node for stock vs ROCm - docs/SAMPLER_CUSTOM_ADVANCED.md: full documentation for agents and users - Category: ROCm Ninodes/Sampling - Bump version to 2.2.0 across all files
204 lines
8.7 KiB
Markdown
204 lines
8.7 KiB
Markdown
# ROCM Ninodes Architecture Documentation
|
||
|
||
## Hardware Configuration
|
||
|
||
### Primary System
|
||
- **Platform**: GMTek Evo-X2 Strix Halo
|
||
- **GPU**: AMD gfx1151 (RDNA 3.5 architecture)
|
||
- **Memory**: 128GB Unified RAM (shared between CPU and GPU)
|
||
- **OS**: Manjaro Linux (kernel 6.17.1-1-MANJARO)
|
||
- **Shell**: Zsh (/usr/bin/zsh)
|
||
|
||
## Software Stack
|
||
|
||
### Core Dependencies
|
||
- **ROCm**: 7.10.0a (latest development build)
|
||
- **PyTorch**: 2.10.0a0+rocm7.10.0a20251011 (nightly build)
|
||
- **Python**: 3.11+ (ComfyUI requirement)
|
||
- **ComfyUI**: Latest main branch with custom optimizations
|
||
|
||
### Environment Configuration
|
||
|
||
#### ComfyUI Launch Flags
|
||
```bash
|
||
--use-pytorch-cross-attention # Enable PyTorch cross-attention for better performance
|
||
--highvram # Optimize for high VRAM systems
|
||
--cache-none # Disable caching to prevent memory issues
|
||
```
|
||
|
||
#### Environment Variables
|
||
```bash
|
||
# ROCm-specific optimizations
|
||
TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 # Enable experimental AOTRITON features
|
||
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True # Enable expandable memory segments
|
||
|
||
# Optional debugging
|
||
ROCM_NINODES_DEBUG=0 # Set to 1 to enable data capture for testing
|
||
```
|
||
|
||
## GPU Architecture Specifics
|
||
|
||
### gfx1151 Optimizations
|
||
- **Precision**: fp16 default (VAE decode is numerically robust; halves memory bandwidth on unified pool)
|
||
- **Tile Sizes**: 768-2048 (larger tiles for high-compression VAEs like LTX)
|
||
- **TF32**: Disabled (not supported/beneficial on gfx1151)
|
||
- **Memory Modifier**: Conservative 1.5x for attention operations
|
||
- **Batch Processing**: Capped at 8 for APU (prevents over-allocation from unified memory reporting)
|
||
- **APU Mode**: Unified memory-aware memory estimation and safety checks
|
||
|
||
### Memory Management Strategy
|
||
- **Attention Memory**: 1.5x modifier for conservative allocation
|
||
- **VAE Processing**: Tiled decoding for large images/videos; minimum tile size floor for high-compression VAEs
|
||
- **Video Support**: 5D tensor processing with 4D output conversion; full-video causal decode for WAN/LTX
|
||
- **VAE Type Detection**: Automatic classification into standard / WAN / LTX Video / pixel-space
|
||
- **Pixel-Space Passthrough**: z-image / z-image-turbo detected and handled as passthrough (no VAE processing)
|
||
- **Cache Management**: Regular `torch.cuda.empty_cache()` calls
|
||
|
||
## Per-Architecture Tuning
|
||
|
||
The VAE decode node dynamically detects GPU architecture and adjusts settings:
|
||
|
||
| Architecture | Family | APU? | Precision | Max Tile | Batch Cap |
|
||
|---|---|---|---|---|---|
|
||
| gfx1150-1151 (Strix Halo) | RDNA 3.5 | ✅ Yes | fp16 | 2048 | 8 |
|
||
| gfx1100-1102 (RX 7900) | RDNA 3 | ❌ No | fp16 | 2048 | 16 |
|
||
| gfx1030-1032 (RX 6900) | RDNA 2 | ❌ No | fp16 | 1536 | 8 |
|
||
| gfx90a (MI200), gfx942 (MI300) | CDNA | ❌ No | bf16 | 2048 | 32 |
|
||
| Other AMD | generic | ❌ No | fp16 | 1024 | 4 |
|
||
|
||
**APU-specific behavior** (Strix Halo):
|
||
- `get_free_memory()` returns system available RAM (not just GPU-reserved)
|
||
- Batch number is capped to 8 (prevents 1.8M batch from 120GB free memory)
|
||
- Memory safety checks use `psutil` for system-wide available memory
|
||
- Precision defaults to fp16 (reduces memory bandwidth pressure on unified pool)
|
||
|
||
## Quantization Support
|
||
|
||
### Supported Quantized Formats
|
||
- **FP8**: Hardware-accelerated FP8 quantization (flux1-dev-fp8.safetensors)
|
||
- **BFloat16**: Native ROCm support with minimal overhead
|
||
- **INT8/INT4**: GGUF format support for WAN 2.2 models
|
||
- **Automatic Detection**: Detects quantized models from filename and dtype
|
||
|
||
### Quantization-Specific Optimizations
|
||
- **Compatibility Mode**: Disables aggressive optimizations for quantized models
|
||
- **Memory Management**: Quantization-aware memory allocation (FP8: 50% vs FP32, INT8: 25% vs FP32)
|
||
- **Dtype Preservation**: Prevents forced dtype conversions that break quantized models
|
||
- **Adaptive Processing**: Smaller tile sizes and chunk sizes for quantized models
|
||
|
||
### User-Reported Issue Fixes
|
||
- **OOM Prevention**: Lower default tile_size (512 vs 768) for compatibility
|
||
- **Batch Optimization**: Disabled by default for quantized models
|
||
- **Memory Cleanup**: Less aggressive cleanup for quantized models
|
||
- **Video Processing**: Adaptive chunk sizing based on frame count and memory
|
||
|
||
### HIPBlas
|
||
- **Status**: Fully supported
|
||
- **Configuration**: Automatic detection and optimization
|
||
- **Performance**: Significant speedup for matrix operations
|
||
|
||
### FlashInfer
|
||
- **Status**: Experimental support for gfx1151
|
||
- **Usage**: Enabled in ROCm builds but may have limitations
|
||
- **Fallback**: Standard attention mechanisms when FlashInfer fails
|
||
|
||
### GPU Direct Storage
|
||
- **Status**: Limited support on current setup
|
||
- **Workaround**: Standard file I/O with memory mapping where possible
|
||
- **Future**: May improve with driver updates
|
||
|
||
## Performance Characteristics
|
||
|
||
### Flux Workflow (1024x1024)
|
||
- **Target Improvement**: 78% over baseline
|
||
- **Checkpoint Load**: <30s target
|
||
- **Sampling**: Optimized for gfx1151 architecture
|
||
- **VAE Decode**: <10s target
|
||
|
||
### WAN Workflow (320x320, 17 frames)
|
||
- **Target Improvement**: 5.6% over baseline
|
||
- **Video Processing**: 5D tensor support with 4D output; causal decode (full video at once)
|
||
- **Memory Usage**: Conservative allocation to prevent OOM
|
||
- **Total Time**: <100s target
|
||
|
||
### LTX Video Workflow (512x512, 49 frames)
|
||
- **VAE Type**: LTX VideoVAE (128 channels, 32x spatial compression, 8x temporal)
|
||
- **Channels**: 128-channel latent → 3-channel RGB output
|
||
- **Precision**: fp16 strongly recommended (128ch × fp32 is memory-prohibitive on APU)
|
||
- **Tile Size**: 1024-2048 recommended (32x compression makes 768→24px latent tiles; floor at 64px)
|
||
- **Processing**: Full-video causal decode (no chunking to avoid temporal artifacts)
|
||
- **Memory**: Memory estimation uses actual channel count (128× that of SD)
|
||
|
||
### z-image / z-image-turbo (pixel-space, any resolution)
|
||
- **VAE Type**: Pixel-space passthrough (latent_channels=3, spatial_compression=1)
|
||
- **Behavior**: No actual VAE decoding — latents are RGB pixels; passthrough in <1ms
|
||
- **Optimizations**: Skips all dtype conversion, model loading, memory estimation, and tiling
|
||
|
||
## Node Architecture
|
||
|
||
### Core Nodes
|
||
1. **ROCMOptimizedCheckpointLoader**: Flux-optimized checkpoint loading
|
||
2. **ROCMOptimizedKSampler**: Basic sampling with ROCm tuning
|
||
3. **ROCMOptimizedKSamplerAdvanced**: Advanced sampling with step control
|
||
4. **ROCMSamplerCustomAdvanced**: V3 drop-in for SamplerCustomAdvanced with ROCm optimizations and enhanced progress tracking
|
||
5. **ROCMSamplerCustomAdvancedBenchmark**: A/B benchmark comparing stock vs ROCm-optimized custom sampler
|
||
6. **ROCMOptimizedVAEDecode**: VAE decode with video support
|
||
7. **ROCMOptimizedVAEDecodeTiled**: Tiled VAE decode for large images
|
||
8. **ROCMVAEPerformanceMonitor**: VAE performance metrics
|
||
9. **ROCMSamplerPerformanceMonitor**: Sampler performance metrics
|
||
10. **ROCMFluxBenchmark**: Flux workflow benchmarking
|
||
11. **ROCMMemoryOptimizer**: Memory management helper
|
||
|
||
### Data Flow
|
||
```
|
||
Input → Checkpoint Loader → KSampler → VAE Decode → Output
|
||
↓ ↓ ↓ ↓
|
||
Monitor Performance Performance Monitor
|
||
```
|
||
|
||
## Testing Infrastructure
|
||
|
||
### Debug Mode
|
||
- **Environment Variable**: `ROCM_NINODES_DEBUG=1`
|
||
- **Data Capture**: Saves input/output tensors to pickle files
|
||
- **Performance Impact**: Zero when disabled
|
||
- **Location**: `test_data/captured/`
|
||
|
||
### Test Categories
|
||
1. **Unit Tests**: Individual node functionality
|
||
2. **Performance Tests**: Timing and memory usage validation
|
||
3. **Integration Tests**: Full workflow execution
|
||
4. **Regression Tests**: Prevent performance degradation
|
||
|
||
## Constraints and Limitations
|
||
|
||
### Hardware Constraints
|
||
- **Memory**: 128GB shared RAM limits batch sizes
|
||
- **GPU**: gfx1151 specific optimizations required
|
||
- **Storage**: No GPU Direct Storage support
|
||
|
||
### Software Constraints
|
||
- **ROCm Version**: Requires 7.10.0a+ for full features
|
||
- **PyTorch**: Nightly builds required for latest optimizations
|
||
- **ComfyUI**: Custom modifications needed for optimal performance
|
||
|
||
### Performance Constraints
|
||
- **Precision**: fp16 default for AMD (VAE decode is inference-only, numerically robust)
|
||
- **Memory**: Conservative allocation to prevent OOM
|
||
- **Batch Size**: Limited by available memory; capped on APU to prevent over-allocation
|
||
- **Video Processing**: 5D tensor handling complexity; causal decode prevents chunking
|
||
|
||
## Future Considerations
|
||
|
||
### Planned Improvements
|
||
- **FlashInfer**: Full support when available
|
||
- **Memory Optimization**: Better memory management strategies
|
||
- **Batch Processing**: Larger batch sizes with memory improvements
|
||
- **Video Support**: Enhanced video processing capabilities
|
||
|
||
### Research Areas
|
||
- **Attention Mechanisms**: Optimize for gfx1151
|
||
- **Memory Layout**: Better tensor organization
|
||
- **Pipeline Optimization**: Reduce data movement overhead
|
||
- **Quantization**: Explore fp16 where stable
|