iGavroche-rocm-ninodes/ARCHITECTURE.md

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

--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

# 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