tt-metal

repository·main·Indexed 23 days ago

https://github.com/tenstorrent/tt-metal

A software stack for Tenstorrent AI hardware featuring TT-NN, a high-level neural network operator library and general compute framework, and TT-Metalium, a low-level kernel programming model.

Tokens
343.1K
Snippets
729
Records
1.4K
Agent score
81%

What's inside tt-metal

  1. Overview of the Blazepose model

    main

    The Blazepose model provided in this repository is an upper-body pose estimation model. It is designed to run completely on the CPU.

    Key characteristics:

    • Output: It outputs 25 keypoints (note that full-body pose estimation typically outputs 33 keypoints).
    • Architecture: The model consists of two distinct networks: a detection network and a landmark network.
    • Source: The implementation is derived from the Blazepose-Pytorch-GitHub source, which ports Mediapipe models from TFLite to PyTorch.
  2. Overview of TT-LLK (Low Level Kernels)

    main

    TT-LLK (Tenstorrent Low Level Kernel library) provides header-only low-level kernels for Tenstorrent AI chips, including Wormhole and Blackhole architectures.

    These kernels act as foundational compute primitives and building blocks for higher-level software stacks that implement machine learning (ML) operations. The library also includes a test environment for validating LLK APIs.

  3. Overview of BGE-Large-EN-v1.5 implementation

    main

    This directory provides the implementation for the BAAI General Embedding (BGE) model bge-large-en-v1.5. The implementation is organized into several components:

    • common.py: Contains model loading utilities specific to BGE.
    • reference/: A PyTorch reference implementation that reuses sentence_bert.
    • ttnn/: The TTNN-based model implementation, also reusing sentence_bert logic.
    • runner/: Infrastructure for running performance benchmarks and tests.

    For implementations specifically optimized for Wormhole devices, refer to the directory models/demos/wormhole/bge_large_en/.

  4. Overview of the Collective Communication Library

    main
    The Tenstorrent Metalium Collective Communication Library provides high-level multi-chip data movement operations. It is designed to work across various deployment scenarios, including N300, T3000, T7000, and T7000 clusters. In this library, a 'configuration' refers to the specific combination of memory layout, data types, kernel arguments (such as dimensions for all-reduce), and allocation strategies used for the operation.
  5. Overview of TT-DiT (Tenstorrent Diffusion Transformers)

    main
    TT-DiT is an implementation of the Diffusion Transformer architecture optimized for parallel execution on Tenstorrent hardware. It is designed for high-performance image and video generation. The architecture is modular, separating core neural network layers, model architectures (Transformers and VAEs), text encoders, and parallelization utilities.
  6. Overview of GPT-OSS-120B Prefill (`gpt_oss_d_p`)

    main

    The gpt_oss_d_p package implements long-context prefill for the GPT-OSS-120B model. It is designed for 4×8 Blackhole Galaxy hardware configurations (TP=8, SP=4, EP=32).

    Key architectural characteristics include:

    • Architecture Style: Built using the "MiniMax-M3 way," reusing the shared DeepSeek EP-MoE dispatch/combine substrate and fused unified_routed_expert_ffn (SwiGLU-OAI + biases).
    • KV Cache: Uses a block-cyclic KV cache.
    • Runtime: Plugs into the model-agnostic models/demos/common/prefill engine.
    • Execution Model: While the KV cache and runtime are chunk-oriented, current implementation (P1-P4) runs single-chunk / one-shot prefill on P1 through P4. Multi-chunk cache-read paths (attending to previous chunks) are planned for P6.
  7. Overview of TT-Metalium

    main

    TT-Metalium is a low-level C++ framework designed for accelerating both ML and non-ML workloads on Tenstorrent hardware. It provides an abstraction between the host system (e.g., x86 CPU) and Tensix hardware, giving programmers full control over hardware and data movement. This allows for highly optimized, tailored kernels for specific use cases.

    TT-Metalium sits at the foundation of the Tenstorrent software stack, below high-level frameworks like TT-Forge/TT-MLIR and the TTNN library.

  8. Overview of TT-Metal

    main
    TT-Metal is a low-level programming model designed for Tenstorrent hardware. It provides user-facing host APIs that allow developers to interact directly with the hardware for high-performance computing tasks. The primary interface for developers is the TT-Metalium host API.
  9. Overview of LLM models for TT-Metalium / TT-NN

    main

    This section provides a performance and compatibility matrix for Large Language Models (LLMs) running on TT-Metalium and TT-NN. It includes details on batch sizes, target hardware (such as QuietBox, Galaxy, and n300), performance metrics like Time to First Token (ttft) and tokens per second (t/s), and the required TT-Metal release versions.

    Key Metrics Definitions:

    • ttft: Time to first token (ms).
    • t/s/u: Tokens per second per user (throughput of the first token generated after prefill).
    • t/s: Total tokens per second (calculated as t/s/u * batch).
    • TP: Tensor Parallelism factor.
    • DP: Data Parallelism factor.

    Performance Notes:

    • Reported LLM performance is based on an input sequence length of 128 (except for Mamba).
    • Performance numbers are collected using tt-metal model demos. Using a vLLM inference server may result in different performance characteristics.
    • Blackhole software optimizations are currently under active development.
  10. Overview of TT-NN

    main

    TT-NN is an open-source C++ and Python library of neural network operations built on top of the TT-Metalium programming model. It provides a PyTorch-like interface for running machine learning workloads on Tenstorrent AI accelerators and serves as the primary high-level API for developing and optimizing ML models.

    Key Capabilities

    • High-Level Neural Network Operations: Optimized implementations of matrix multiplication, convolution, attention, data movement, collective communications (CCLs), element-wise ops, reductions, losses, and pooling.
    • Tensor Library: A flexible abstraction for managing multidimensional arrays across host and device, allowing precise control over data layout across clusters.
    • Native Multi-Device Support: Virtualizes multiple Tenstorrent devices into a single logical unit for scaling across clusters.
  11. Overview of PatchTSMixer for TT-NN

    main

    PatchTSMixer is a lightweight MLP-Mixer based architecture designed for multivariate time series forecasting. It is implemented using TTNN APIs specifically for Tenstorrent Wormhole hardware.

    Architecture Flow: Patchify $\rightarrow$ Linear Projection $\rightarrow$ Positional Encoding $\rightarrow$ Mixer Blocks $\rightarrow$ Forecast Head.

    Mixer Block Components:

    • Patch Mixer (temporal)
    • Feature Mixer (hidden)
    • Optional Channel Mixer (cross-variate)

    I/O Shapes:

    • Input: (batch, context_length, num_channels)
    • Output: (batch, prediction_length, num_channels)