TT-Forge AI Compiler Stack

The tt-lang-models repository contains several high-performance implementations of various model architectures optimized for Tenstorrent hardware using TT-Lang. Key models include:

• DFlash: A lightweight cross-attention draft model for speculative decoding. It uses TT-Lang kernels (RoPE, RMSNorm, SiLU, residual adds) to achieve significant decoding speedups.
• Engram: A port of the DeepSeek Engram conditional memory module to TT-Lang on Wormhole, utilizing streaming dataflow kernels and PipeNet for inter-core boundary sharing.
• nanochat: An inference implementation written entirely in TT-Lang.
• Oasis: Real-time world generation using the Oasis 500M diffusion transformer, running end-to-end inference in a single captured trace.
• Qwen-Image: A TTNN + TT-Lang implementation of Qwen-Image 20B image generation.
• Micelle MD: Cell-list molecular dynamics implementation featuring Ewald electrostatics and on-device Verlet integration.
• Toy World Model: A Pong world model based on a diffusion transformer running on a single Blackhole card.

TT-NN supports multi-device operations using a Single-Program Multiple-Device (SPMD) technique. This allows you to scale single-device application code to multiple connected devices by parallelizing computations across them.

To scale models, TT-NN uses three core abstractions:

1. MeshDevice: A "virtual device" abstraction that defines a logical 2-D mesh of connected physical devices. Operations targeting a MeshDevice are distributed via SPMD across all devices in that mesh.
2. Input Data Distribution: The mechanism for distributing host-memory data to the on-device memory of the MeshDevice. When an operation is distributed, the code within a single-device scope operates only on its local shard of the data.
3. MeshTensor: A collection of tensor shards distributed across the 2D mesh. While a standard Tensor is an N-dimensional matrix, a MeshTensor represents the global view of data spread across the MeshDevice.

A MeshDevice is a virtual device abstraction that bundles multiple physical devices into a single logical unit to enable parallel execution across a mesh topology. It is natively supported at the runtime level to reduce dispatch overhead via unified command queues.

Key Optimizations:

• Kernel Compilation Broadcasting: Compilation artifacts are automatically broadcast to all devices in the mesh to avoid redundant per-device compilation.
• Data Broadcasting: Replicated tensors are efficiently broadcast across the mesh topology with minimal host-side processing.
• Unified Command Dispatch: Uses mesh-aware command queues to coordinate lock-step parallel execution.

Important Note: MeshDevice does not hide the distributed nature of computation. Users must explicitly manage:

• Data Distribution: Specifying how tensors are sharded or replicated using mesh mappers.
• Collective Communication: Using CCL operations (e.g., all-gather, reduce-scatter) for device communication.

import ttnn

# Example: Instantiating a MeshDevice over a 1x2 Wormhole configuration
mesh_device = ttnn.open_mesh_device(ttnn.MeshShape(1, 2))

Metal Trace is a performance optimization feature designed to eliminate host overhead involved in constructing and dispatching operations. It is most effective when the host time required to dispatch a model exceeds the device time needed to execute the operations (i.e., the model is host-bound).

How it works: Metal Trace records the commands for dispatching operations into a DRAM buffer. During execution, these recorded commands are replayed from the buffer.

Key Constraints:

• Static Parameters: All operation parameters, including input/output tensor shapes and memory addresses, must be static. It cannot be used directly for generative models with changing sequence lengths without advanced techniques like multiple traces.
• Ideal Use Cases: Statically sized inputs/outputs, such as image or sequence classification.
• Requirements: You must use program cache and ensure target operations are already compiled before capturing the trace.

A grid is the space of nodes (Tensix Cores) where a kernel is executed. In single-chip mode, the grid is 2D. In multi-chip mode, it includes additional dimensions for connectivity (e.g., same card, same host).

Use ttl.grid_size(dims) to get the total size of the grid in a specific dimensionality. If the requested dims is smaller than the actual grid rank, the highest rank dimension is flattened. If dims is larger, higher ranks are padded with 1.

Use ttl.node(dims) to get the zero-based coordinates of the current node. If dims is smaller than the grid rank, the highest rank is flattened. If dims is larger, higher ranks are padded with 0.

# Grid size examples
# For (8, 8) single chip or SPMD grid:
x_size = ttl.grid_size(dims=1)          # returns 64
x_size, y_size = ttl.grid_size(dims=2) # returns (8, 8)
x_size, y_size, z_size = ttl.grid_size(dims=3) # returns (8, 8, 1)

# Node coordinate examples
# For (8, 8) single chip or SPMD grid:
x = ttl.node(dims=1)          # returns x in [0, 64)
x, y = ttl.node(dims=2)       # returns (x, y) in [0, 8) x [0, 8)
x, y, z = ttl.node(dims=3)    # returns (x, y, 0) in [0, 8) x [0, 8) x [0]