code/models/tt_dit/README.md

TT-DiT: Tenstorrent Diffusion Transformers

This directory contains the implementation of the Tenstorrent Diffusion Transformer (TT-DiT) architecture, designed for optimized parallelism of DiTs for image and video generation.

Supported Models

For detailed information about each model including performance metrics, usage instructions, and specific requirements, see:

• Stable Diffusion 3.5 Large - Text-to-image generation
• Flux 1 - Text-to-image generation (schnell & dev variants)
• Motif - Text-to-image generation model
• Qwen-Image - Text-to-image generation model
• Ideogram 4.0 - Text-to-image generation model (best-in-class typography; Blackhole)
• Mochi-1 - Video generation model
• Wan2.2-T2V-A14B - Text-to-video generation model
• See experimental/ for in-progress model variants (AniSora, Wan2.2-Distill, LoRA adapters)
• LTX-2.3 - Joint audio-video generation model (Pro and Fast variants)

Directory Structure

tt_dit/
├── layers/              # Core neural network layers
├── models/              # Model architectures and documentation
│   ├── transformers/    # Transformer implementations (SD35, Mochi, Wan, LTX, Flux1, Motif, QwenImage)
│   ├── vae/            # VAE/Autoencoder implementations
│   ├── StableDiffusion35.md  # SD3.5 model documentation
│   ├── Flux1.md         # Flux 1 model documentation
│   ├── Motif.md         # Motif model documentation
│   ├── QwenImage.md     # Qwen-Image model documentation
│   ├── Mochi_1.md       # Mochi-1 model documentation
│   ├── Wan2_2.md        # Wan2.2 model documentation
│   └── LTX2.md          # LTX-2.3 model documentation
├── encoders/            # Text encoder implementations
│   ├── clip/           # CLIP encoder
│   ├── gemma/          # Gemma encoder
│   └── t5/             # T5 encoder
├── parallel/            # Parallelization utilities
│   ├── config.py        # Parallel configuration
│   └── manager.py       # Parallel execution management
├── pipelines/           # End-to-end model pipelines
│   ├── stable_diffusion_35_large/
│   ├── mochi/
│   ├── wan/
│   ├── ltx/
│   ├── flux1/
│   ├── motif/
│   └── qwenimage/
├── tests/              # Test suite
│   ├── models/         # Model-level tests (sd35, mochi, wan2_2, ltx, flux1, motif, qwenimage)
│   ├── encoders/       # Encoder tests
│   ├── blocks/         # Block-level tests
│   └── unit/          # Unit tests for layers
└── utils/             # Utility functions
    ├── check.py       # Validation utilities
    ├── padding.py     # Padding operations
    ├── substate.py    # State management
    └── tensor.py      # Tensor operations

Core Components

Layers

• Embeddings: Position and token embedding implementations
• Feedforward: MLP and feed-forward network components
• Linear: Optimized linear transformation layers
• Normalization: Layer normalization implementations

Models

• Transformers: DiT transformer implementations for various generative models
  • Support for multiple architectures including SD3.5, Mochi, Wan2.2, LTX-2, Flux1, Motif, and Qwen-Image
  • Model-specific attention mechanisms and transformer architectures
• Autoencoders: VAE implementations for different models

Parallel Processing

• Config: Parallel configuration management with DiTParallelConfig
• Manager: Parallel execution and device management

Pipelines

End-to-end pipeline implementations for multiple generative models:

• Stable Diffusion 3.5 Large: Text-to-image generation (1024x1024px)
• Flux 1: Text-to-image generation (schnell & dev variants, 1024x1024px)
• Motif: Text-to-image generation (6B model, 1024x1024px)
• Qwen-Image: Text-to-image generation (1024x1024px)
• Mochi-1: Video generation model (824x480px, 168 frames)
• Wan2.2-T2V-A14B: Text-to-video generation
• LTX-2.3: Joint audio-video generation (Pro one-stage and Fast distilled two-stage)

Each pipeline includes:

• Automatic parallel configuration for different device meshes
• Optimized execution on Wormhole systems
• Comprehensive timing collection and profiling

Testing

The test suite is organized into three main categories:

1. Unit Tests (tests/unit/):

   • Layer-level testing
   • Individual component validation
   • Coverage for embeddings, feedforward, linear, and normalization layers

2. Model Tests (tests/models/):

   • Smoke tests for model components
   • Performance testing
   • Pipeline validation

Running tests:

# Run tests from tt-metal/

# Run unit tests
python -m pytest  models/tt_dit/tests/unit/

# Run all model tests
python -m pytest  models/tt_dit/tests/models/

# Run specific model pipeline tests
python -m pytest models/tt_dit/tests/models/sd35/test_pipeline_sd35.py -v
python -m pytest models/tt_dit/tests/models/flux1/test_pipeline_flux1.py -v
python -m pytest models/tt_dit/tests/models/motif/test_pipeline_motif.py -v
python -m pytest models/tt_dit/tests/models/qwenimage/test_pipeline_qwenimage.py -v
python -m pytest models/tt_dit/tests/models/mochi/test_pipeline_mochi.py -v
python -m pytest models/tt_dit/tests/models/wan2_2/test_pipeline_wan.py -v
python -m pytest models/tt_dit/tests/models/ltx/test_pipeline_ltx_distilled.py -v
python -m pytest models/tt_dit/tests/models/ltx/test_pipeline_ltx_two_stages.py -v

Key Features

• Modular Architecture: Clean separation between layers, models, and parallel processing
• Optimized Parallelism:
  • Automatic configuration of tensor and sequence parallelism
  • Efficient device mesh management
  • Dedicated parallel managers for different model components
• Performance Monitoring:
  • Built-in timing collection
  • Performance profiling support
  • Tracing capabilities for optimization

Contributing

When adding new components:

1. Place implementation in appropriate subdirectory
2. Add corresponding unit tests
3. Update model tests if necessary
4. Ensure parallel processing compatibility