TT-Metaliumによるカーネルプログラミング

Tenstorrent Tech Talk #2, Session2

Transcript

1. Tenstorrent Tech Talk #2 TT-Metaliumによるカーネルプログラミング July 16 2025 Kohei Yamaguchi / Tenstorrent Japan @sott0n
2. Tensix Core Architecture Compute (Tensix Engine) RISC-V 2 RISC-V 3 RISC-V 4 RISC-V 5 RISC-V 1 Router 0 L1 Memory • 5 “Baby RISC-V” Cores (32-bit RISC-V ISA) • 2 Network-on-Chip (Router) • 1.5MB SRAM Cache (L1) Compute (TensixEngine) Vector Math Engine • Tile/Matrix Math Engine (FPU) • Vector Math Engine (SFPU) Router 1 Tile/Matrix Math Engine PCIe Ethernet GDDR6 GDDR6 GDDR6 GDDR6 GDDR6 GDDR6 GDDR6 GDDR6 Blackhole C T D E P A RISC-V CPUs Tensix cores DRAM cores ETH cores PCIe core ARC core
3. Tenstorrent OSS Stack Partners TT-Forge vLLM Python kernels TT-NN TT-Metalium TT-LLK (low-level-kernels) PyTorch models TT-Fabric (unified scale-up and scale-out) Manually optimized models TT-Train TT- Transformer LLM training LLM inference models LLM, t2s, s2t models Jax PyTorch TF ONNX Models AI Workloads Open Source Partners Tenstorrent Open Source Software
4. TT-Metalium: Built for AI and Scale-Out Native Multi-Device Kernels and Ops TT-Metalium GPU Programming Deep Learning Ops Collective Comms Ops VS. TT-Metalium C++ Host API TT-Metalium C++ Kernel API TT-NN C++ Host API GPU Kernel Language DNN CCL
5. Programming Model Host Side
6. Host-Device Programming Model | Bird’s Eye View Device Host Program Command Queue PCIE Device object Core Mesh SRAM RISC-V 1 RISC-V 2, 3, 4 RISC-V 5 Data movement kernel Router 0 Matrix and vector engines Router 1 Compute kernel Data movement kernel Single Tensix Core
7. TT-Metalium Command Queue Host Machine Main Configuration Program Command Queue To Device To Host Device Metal Runtime Support (program execution) 1. EnqueueWriteBuffer() ホストからデバイスへデータを転送 （device DRAMへ書き込む） 2. Send commands デバイスはCQからのコマン ド列を待ち、受け取ったあと 実行する 3. Requests output data ホストがデバイスのoutput dataを要求する 4. EnqueueReadBuffer() デバイスからホストへデータを転送 （device DRAMから読み込む）
8. Main Configuration Program TT-Metalium Host Program Organization Initialize the Program and Device: • device object の作成 • program objectを作成 • core coordinatesの指定 Initialize Program and Device Configure and Create Collaboration Mechanisms Configure and Create Kernels Dispatch Program to Device for Execution Finish Program Execution Allocate Device Memory and Initialize Source Data Allocate Device Memory: • DRAMとSRAMのメモリレイアウトを設定 (sharded or interleaved) • DRAM and SRAM buffer の作成 Initialize Source Data: • Source data用のベクトルを作成（buffers） • 作成したbuffersにSource dataを書き込み
9. Main Configuration Program TT-Metalium Host Program Organization Configure and Create Collaboration Mechanisms: • Command Queue (CQ) object の作成 • Circular Buffer (CB) layout の設定と、データ転送のためのCB objects の作成 • (if necessary) Atomic Counter (AC) objects の作成 Initialize Program and Device Allocate Device Memory and Initialize Source Data Configure and Create Collaboration Mechanisms Configure and Create Kernels Dispatch Program to Device for Execution Finish Program Execution Configure and Create Kernels Configure and Create Kernels: • Kernels の作成 • Compute kernel • Data movement kernel x2 • Kernel ごとに使用するコアを指定するなどの設定
10. TT-Metalium Host Program Organization Dispatch the Program to Device for Execution: • 各kernel のRuntime引数を指定 • EnqueueProgram() がデバイスにプログラムを書き込み、実行する Main Configuration Program Initialize Program and Device Allocate Device Memory and Initialize Source Data Configure and Create Collaboration Mechanisms Configure and Create Kernels Dispatch Program to Device for Execution Finish Program Execution Finish Program Execution Finish Program Execution: • Finish() は、指定された CQ objectを通して、ディスパッチされたすべてのコマ ンドが完了するまでブロックする • EnqueueReadBuffer() は、デバイスの出力データをバッファからホストメモリ に読み込む • 最後に、デバイスをCloseする
11. Host code: Part 1 } 指定したdevice_idのdeviceを確保 } Command queueとProgramのobjectを作成 Core coordinates, tile size, number of tiles, buffer size SRAM上のsrc0 CB の定義 DRAM buffer のConfig設定 Source/Destination のDRAM buffer の作成 SRAM上のsrc1 CB の定義 SRAM上のoutput CB の定義
12. Host code: Part 2 データ転送のReader kernal の作成 Randomなベクトルデータを作成し、 ホストからsrc1 DRAM buffer へデータ転送 Constantなベクトルデータを作成し、 ホストからsrc2 DRAM buffer へデータ転送 各kernel にRuntime argumentsを設定 Kernelを実行し、処理が完了するまで待機 結果をデバイスからホストへ転送 データ転送のWriter kernal の作成 Compute kernel の作成 - Math Fidelityなど計算時の条件を設定
13. Programming Model Device Side
14. Tensix Core Subsystems • Communication using NoC: • NoC (Network-on-Chip)は、RISC-V Coreを搭載し、 デバイス・コンポーネント間でデータを移動する • 5つのBaby RISC-V Coreのうち、両サイドの2つの RISC-V Core 1/5でNoCの制御を行う • Compute: • 中央の3つのRISC-V Coresで命令やデータ転送を制 御する • Tensix Engineとして、Matrix Engine (FPU)と Vector Engine (SFPU) などがある • Memory subsystem: • SRAM (L1) • DRAM Com e ( ensi Engine RISC V 2 RISC V 3 RISC V RISC V 5 RISC V Ro er Memory Ro er
15. Communication using NoC of a Tensix Core • Baby RISC-V core for NoC: • NoCによるデータ通信はRISC-V 1とRISC-V5のCoreを用いる • ユーザが定義したデータ転送Kernelが実行される • RISC-V 1 は主にReader Kernelとして振る舞う • RISC-V 5 は主にWriter Kernelとして振る舞う • これらのBaby RISC-V Coreは単純な計算能力を持つ • この仕組みによってSRAMからコードを実行し、NoCを介して データの通信を行う On-chip SRAM (L1) RISC-V 1 Router 0 RISC-V 5 Router 1 NOC instructions Dispatching データ転送 Kernal (Reader/Writer) を定義して、NoCを通じてデータ転送を行う
16. Computation Subsystem of a Tensix Core RISC-V 2によるデータロード: Unpacker • SRAMからデータをFetch → データ変換 → Sourceレジスタへのデータロード （SrcA/SrcB） RISC-V 3による計算実行: Matrix, Vector Engine, etc • Math Engineに対して命令を発行する • Matrix Engineは32x32を1 tileとしてMatmulやdot product命令などを実行す る（Tensix Coreでは最も高性能なエンジン） • Vector Engineは、GELUやExp,Sqrtなどをサポート RISC-V 4によるデータストア: Packer • Destinationレジスタからデータを取得 → データ変換 → SRAMへデータをス トア` Compute Kernelによって、命令・データのRead、計算実行、結果のStoreという一連の計算を行う On-chip SRAM (L1) RISC-V 2 Unpacker RISC-V 3 Matrix and vector engines RISC-V 4 Packer Compute instructions dispatching Src/Dest register
17. Reader/Writer, Compute Kernels Compute RISC-V 2 RISC-V 3 RISC-V 4 RISC-V 5 RISC-V 1 Router 1 Router 0 L1 Memory NoC 0 Reader kernel Writer kernel Compute kernel NoC 1 CBs CBs Reader kernel Compute kernel Writer kernel
18. Register layout and Data flow for Matrix/Vector Engine Matrix Engine (FPU) Vector Engine (SFPU)
19. Circular Buffers (CB) はSRAMのアドレスと入力と出力に 関する情報を含むデータ構造 CBを経由したDataflow kernel と Compute kernel の同期 • dest regの半分を使って信号を送る同期呼び出し • 入力データに対する1タイル分のAdd演算を実行され、dest regに格納 • dest regの結果がSRAMに転送 • dest regの解放 CBに対してタイル分をpop/pushすることで、CBのスペースを開放/確保する これによって、Dataflow kernelとCompute kernel間で同期を取る Device code: Elementwise Add case (compute kernel)
20. Compile Time and Runtime
21. TT-Metalium Program Compilation Input: • Source code: • TT Deviceを操作するためのTT- Metalium APIを使って書かれた プログラム • TT-Metalium API Library: • TT Deviceに対する操作のため の以下のようなAPIを持ったラ イブラリ • Data movement APIs • Circular buffer APIs • Compute APIs • Kernel argument APIs Output: Program object: Host binary code、Kernel Object、Data movement Object を含んだTT Deviceターゲットの Object Host Machine TT-Metalium TT-Metalium API library Source code Compile Program object Host binary code Compute kernel binary code Data movement kernel binary code
22. TT-Metalium Runtime Components 1. MemoryのAllocationとInitialization 2. TT DeviceへのKernels発行 3. Kernelsに対するRuntime argumentsの設定 4. TT Device上での演算の開始 5. TT Device上の実行結果の回収
23. Host Machine TT-Metalium Program object Host binary code Compute kernel binary code Data movement kernel binary code TT-Metalium Runtime | Allocate and initialize memory Device Tensix Core 1 On-chip SRAM Buffer DRAM Core Off-chip DRAM Buffer Tensix Core 2 On-chip SRAM Buffer Hostは、実行時に DRAM bufferやCircular bufferなどのBufferを指定されたコアに作成する
24. Host Machine TT-Metalium Program object Host binary code Compute kernel binary code Data movement kernel binary code TT-Metalium Runtime | Set kernel runtime arguments Host は、KernelのRuntime引数を設定する Device Tensix Core 1 On-chip SRAM DRAM Core On-chip SRAM Tensix Core 2 On-chip SRAM Kernel Runtime arguments Kernel Runtime arguments Kernel Runtime arguments
25. Host Machine TT-Metalium Program object Host binary code Compute kernel binary code Data movement kernel binary code TT-Metalium Runtime | Program execution on device Off-chip DRAM Core Mesh PCIe Device Management and Control Kernel Programの処理が終了したら、 TT Device上にある結果をHostへ転送する TT-Metalium Runtimeは、KernelをTensix Coreに送信し、 Kernel Programの処理が終了するまで管理をする
26. Metalium Execution Model
27. TT-Metalium Execution Model • Create/Compile Program • DRAM Bufferの作成 • DRAM bufferとSRAM上のCB間でデータのやり取り • SRAMにCircular Buffer (CB) の作成 • CBsを用いてSignalを送るイメージでKernel間のデータの同期を行う • Kernelsの作成 • Data Movement Kernel (Baby1) • DRAMからSRAMへのデータロード • Compute Kernel (Baby2/3/4) • Matrix (FPU), Vector (SFPU)のMath Engineによる計算 • Data Movement Kernel (Baby5) • SRAMからDRAMへのデータストア • Load program and runtime config • E.g. DRAM buffer start address, size, coordination (3, 4) 1. Create, compile, load Program 2. Create DRAM buffer 3. Load data to DRAM buffer 4. Launch Program 5. Store data from DRAM buffer Stream Src0 RISC-V 1 Src1 Out RISC-V 2 RISC-V 3 RISC-V 4 RISC-V 5 Compute Stream NoC NoC CB Src0 CB Src1 Program CB Out Host DRAM Tensix Core SRAM
28. Kernel and Memory Model Reader Kernel Compute Kernel Writer Kernel Tensix Core Grid Global 8-bank DRAM • Circular Buffer(CB) は各コアのSRAM上に存在 • L1は全てのTensixCoreがNoC経由でアクセスできる共有SRAM • プログラムはL1に格納され、スタック上のContentを保持するための Private Memoryがコアごとに存在する Buffer object はDRAMの連続的なメモリ空間を抽象化したもの ページ単位で複数のBankに分散配置される CB0 CB1 CB2 CB31 SRAM (L1) Tensix Core RISC-V private memory Tensix local memory
29. Tensor Layout Page 0 = [0,0 to 0,63] Page 1 = [1,0 to 1,63] Page 2 = [2,0 to 2,63] Page 62 = [62,0 to 62,63] Page 63 = [63,0 to 63,63] Page 0 Page 1 Page 2 Page 3 Row-Major Layout Tiled Layout 1 Tileが1 pageとして保存される 64x64のTensorの場合、32x32の4 Tiles(4 pages)になる 2D Tensorの各行が1 pageとして保存される 64x64のTensorの場合、Bufferの中に64 pages存在する
30. Memory Layout: Interleaved Interleaved layout Tensorを構成する複数のPageが各DRAM bankにラウンドロビンで保存される ただし、新しいTensorの割り当ては常に最初のbankからスタートするためfragmentationが発生する可能性がある
31. Memory Layout: Shared Shared layout Tensorを構成する複数のPageを各コアのSRAM上に分散させて保存する • コアのローカルメモリが活用できるため、DRAMにアクセスするInterleavedよりも高速 • 計算パターンにTensorのデータ分散戦略を合わせるとデータ通信の距離によるオーバーヘッドが短縮される
32. Native Tile based computing • Metaliumでは, 32x32のTile をNativeにサポートしており、TT HardwareはこのTile単位での 行列計算を前提に動作するように組み込まれている • 48x1024のMatrixの計算では、Paddingによって64x1024にアラインメントされて計算する
33. Programming Matmul
34. Matrix Multiply - Matmul • Multiply 2 matrices: C = A x B • A is M x K • B is K x N • C is M x N A B K M N K X C = M N
35. Host code: Matmul M,K,Nに対する Tile数の計算 DRAMの設定 • Sizeの指定 • Buffer typeの指定 など DRAM Bufferの作成 CBの設定 • Page sizeの設定 • CBの長さの設定 • DataFormatの指定 など BFloat16 & HiFi4 で設定 Kernelsの設定 CompileTime引数と して、Mt/Mt/Ntを Matmul kerneに与え る Single tileのサイズ Runtime引数の設定 Command Queue へのPush
36. Reader Kernel: Matmul • InterleavedAddrGenFast • NoC座標を指定せずにDRAMの指定のBankへアクセスできる • DRAMはラウンドロビンでデータが格納されているため、コンパイ ル時にNoC座標はわかっている • Mt, Nt, Ktを基準にしてデータをRead • Mt, Nt, Kt: Host側で計算した各次元のTile数 • 同期機構を用いたRead • noc_async_read_tileとnoc_async_read_barrier の命令に より、BarrierをしつつNoC経由でDRAMからデータをRead • Readが完了したら、cb_push_backで1 tile分Pushして、 Compute kernel にReadが完了したことを伝える
37. Compute Kernel: Matmul • FPUの初期化 • Circular Bufferのin0/1とoutを設定 • Tile単位でMatmulの計算 • Circular Buffer経由でReader kernel と同期し、データが SRAMに書き込まれるまで待機 • TT-LLK APIのmatmul_tiles でMatmulを計算 • Circular Bufferにinput tileを消費したとMarkする • 計算結果をPackしてSRAMにWrite • Packer によりdest register にある計算結果をCBに書き込む • 書き込みが完了したら、output tileがReady状態になったと Markする
38. Writer Kernel: Matmul • 基本的にReaderと同様の処理を記述 • Matmul: C=A*B の Cへの出力 • InterleavedAddrGenFast により、DRAM bankへのアドレスを自動で計算 • noc_async_* 命令を使って、NoC経由でDRAMへのWrite • cb_* 命令を使って、Compute Kernelとの同期
39. Matmul: Single Core Compute UNPACK MATH Matrix/Vector PACK Router 1 Router 0 L1 Memory NoC 0 NoC 1 CBs CBs noc_async_read noc_async_read_barrier noc_async_write noc_async_write_barrier unpack_AB cb_push_back for input pack_C cb_pop_front for dest cb_pop_front for input cb_push_back for dest Wait tiles Wait tiles Wait tiles Wait tiles tt-llk: Matmul • Create/Close Device • Input Vector • data format • data layout • Tilize/Untilize • ---------- • Create Program • DRAM/Circular Buffers • Create Kernels • Read/Write/Compute • Compile time args • Runtime args • Enqueue Single Core
40. Matmul optimization: Multi Core + Reuse • Create/Close Device • Input Vector • data format • data layout • Tilize/Untilize • ---------- • Create Program • DRAM/Circular Buffers • Create Kernels • Read/Write/Compute • Compile time args • Runtime args • Enqueue • Data Blocks/Subblocks • Reuse same blocks • on cores row/column • split cores • Same compute kernel • R/W kernels • Matrix chops Multi Core Multi Cores + Reuse Multi Core (DRAM to each core) Single Core T T T T T T T T T D (0,0) (1,0) (2,0) (0,1) (0,2) X = Reuse, Block/Subblocks
41. Matmul optimization: Multi-cast • Create/Close Device • Input Vector • data format • data layout • Tilize/Untilize • ---------- • Create Program • DRAM/Circular Buffers • Create Kernels • Read/Write/Compute • Compile time args • Runtime args • Enqueue Single Core • Data Blocks/Subblocks • Reuse same blocks on cores row/column • split cores • Same compute kernel • R/W kernels - Matrix chops Multi Core Multi Cores + Reuse • Multi-cast - DRAM to 1st row/col - 1st row/col to rest Multi Cores + Reuse+ Multicast • Padding • Mixed-Precision • Fused Ops • 2D/1D mcasting • Intermediates on L1 More optimization Multi-cast T T T T T T T T T D (0,0) (1,0) (2,0) (0,1) (0,2) X = DRAMからは1行目, 1列目のコアにデータを転送 他のコアへのデータ転送はコアからコアへ
42. Reference • TT-Metalium documentation • TT Architecture and Metalium Guide • TT-Metalium Tech Reports • Programming Examples • Matmul Example • Matmul SingleCore Example • Matmul MultiCore Example • Matmul MultiCore Optimized Example • tt-isa-documentation • Memory on Tenstorrent by Martin • Tenstorrent Wormhole Series Part7: Bits of the Matmul by Corsix
43. Thank you