உருவாக்கப்பட்ட Triton குறியீடு, GEMM எடுத்துக்காட்டுகள் மற்றும் ROCm warp-size பிரச்சினைகள்

இந்த மூன்று நடைமுறை வழிமுறைகளை ஆராய்வது, அதாவது torch.compile-ஆல் உருவாக்கப்பட்ட Triton குறியீட்டை பரிசீலிப்பது, tile செய்யப்பட்ட மாத்திரைப் பெருக்கல் (matrix multiplication, GEMM) குறியீட்டை ஒப்பிடுவது, மற்றும் CUDA-ஐ AMD ROCm-க்கு மாற்றும்போது எழும் warp-size பிரச்சினைகளைக் கையாள்வது, உயர்மட்ட AI framework-கள் எவ்வாறு நேரடியாக hardware அறிவுறுத்தல்களாக மொழிபெயர்க்கப்படுகின்றன என்பதை விளக்குகிறது.

1. torch.compile ஆல் உருவாக்கப்பட்ட Triton குறியீட்டைப் பரிசீலித்தல்

RMSNorm (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * weight) போன்ற ஒரு வழக்கமான PyTorch செயல்பாட்டை @torch.compile மூலம் மடக்கும்போது, TorchInductor பின்னணியில் ஒரே ஒரு இணைக்கப்பட்ட (fused) Triton kernel-ஐ உருவாக்குவதன் மூலம் நான்கு தனித்த HBM round-trips-ஐ நீக்குகிறது.

TORCH_LOGS="output_code" python script.py என உங்கள் script-ஐ இயக்கினால், compiler-ஆல் உருவாக்கப்பட்ட உண்மையான Triton குறியீடு அச்சிடப்படும்; அது இவ்வாறு இருக்கும்:

import triton
import triton.language as tl

@triton.jit
def triton_per_fused_mul_pow_rsqrt_0(in_ptr0, in_ptr1, out_ptr0, xnumel, rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4096
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rindex = tl.arange(0, RBLOCK)[None, :]
    rmask = rindex < rnumel

    # 1. Single coalesced read of input tensor (x) and weight vector from HBM
    x = tl.load(in_ptr0 + (rindex + (4096 * xindex)), rmask & xmask, other=0.0).to(tl.float32)
    w = tl.load(in_ptr1 + (rindex), rmask, eviction_policy='evict_last').to(tl.float32)

    # 2. Compute square, row-wise mean, and reciprocal square root entirely in registers
    x_sq = x * x
    row_sum = tl.sum(x_sq, 1)[:, None]
    rstd = tl.rsqrt((row_sum / 4096.0) + 1e-06)

    # 3. Multiply by inverse standard deviation and learned weight, then write once
    out = (x * rstd) * w
    tl.store(out_ptr0 + (rindex + (4096 * xindex)), out, rmask & xmask)

• Compiler தானாகக் கண்டுபிடித்தது: weight vector w-இல் உள்ள eviction_policy='evict_last' என்பதைக் கவனியுங்கள். w ஒவ்வொரு வரிசையிலும் (xindex) மீண்டும் மீண்டும் பயன்படுத்தப்படுகிறது என்பதை TorchInductor தானாகவே கண்டறிந்து, x-ஐ ஒரே ஒரு முறை streaming செய்யும்போது w-ஐ chip-மீது பிணைத்து வைக்குமாறு GPU-யின் L2 cache-க்கு அறிவுறுத்துகிறது.

2. அடுத்தடுத்து ஒப்பிடக்கூடிய மாத்திரைப் பெருக்கல் (Matrix Multiplication, GEMM) குறியீடு: CUDA C++ எதிர் Triton

மாத்திரைப் பெருக்கல் (C = A × B), CUDA C++-இல் உள்ள கைமுறை thread synchronization-க்கும் Triton-இல் உள்ள declarative tile கணிதத்திற்கும் இடையிலான பெரும் வேறுபாட்டைத் தெளிவுபடுத்துகிறது.

A. CUDA C++ (வெளிப்படையான Shared Memory Tiling & Sync Barriers)

CUDA C++-இல், ஒவ்வொரு thread-ம் ஒரு scalar வெளியீட்டு உறுப்பை (acc) கணக்கிடுகிறது; அதே வேளை thread-கள் tiles-ஐ __shared__ memory-யில் ஏற்ற கைமுறையாக ஒத்துழைக்கின்றன:

#define TILE 16

__global__ void gemm_cuda(const float* A, const float* B, float* C, int M, int N, int K) {
    __shared__ float sA[TILE][TILE];
    __shared__ float sB[TILE][TILE];

    int row = blockIdx.y * TILE + threadIdx.y;
    int col = blockIdx.x * TILE + threadIdx.x;
    float acc = 0.0f;

    for (int t = 0; t < (K + TILE - 1) / TILE; ++t) {
        // 1. Each thread manually loads 1 element into Shared Memory
        sA[threadIdx.y][threadIdx.x] = (row < M && t * TILE + threadIdx.x < K)
                                       ? A[row * K + t * TILE + threadIdx.x] : 0.0f;
        sB[threadIdx.y][threadIdx.x] = (t * TILE + threadIdx.y < K && col < N)
                                       ? B[(t * TILE + threadIdx.y) * N + col] : 0.0f;
        __syncthreads(); // Mandatory barrier before reading shared tiles

        // 2. Inner dot product across the 16x16 tile
        for (int k = 0; k < TILE; ++k) {
            acc += sA[threadIdx.y][k] * sB[k][threadIdx.x];
        }
        __syncthreads(); // Mandatory barrier before loading next tile
    }

    if (row < M && col < N) C[row * N + col] = acc;
}

B. OpenAI Triton (Declarative 2D Tile Accumulation & Tensor Cores)

Triton-இல், ஒரே ஒரு program block, tl.dot() பயன்படுத்தி முழு BLOCK_M × BLOCK_N வெளியீட்டு tile-ஐயும் கணக்கிடுகிறது; இது தானாகவே hardware Tensor Cores-ஆகத் தாழ்த்தப்படுகிறது:

@triton.jit
def gemm_triton(A_ptr, B_ptr, C_ptr, M, N, K,
                stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
                BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)

    a_ptrs = A_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = B_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)

    # Accumulate in FP32 registers while multiplying FP16 tiles on Tensor Cores
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, K, BLOCK_K):
        a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (k + offs_k[None, :] < K), other=0.0)
        b = tl.load(b_ptrs, mask=(k + offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0)
        acc = tl.dot(a, b, acc)  # Hardware Tensor Core instruction
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    c_ptrs = C_ptr + stride_cm * offs_m[:, None] + stride_cn * offs_n[None, :]
    tl.store(c_ptrs, acc.to(tl.float16), mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))

3. CUDA C++-ஐ AMD ROCm (HIP)-க்கு மாற்றும்போது எழும் Warp-Size பிரச்சினைகள்

AMD-யின் HIP interface வழியாக தாழ்மட்ட CUDA C++ kernels-ஐ AMD Instinct GPU-களுக்கு (MI300X அல்லது MI325X போன்றவை) மாற்றும்போது, syntax மொழிபெயர்ப்பு (cudaMalloc → hipMalloc) எளிதானது. எனினும், NVIDIA warp-களில் 32 threads இருக்கும்போது, AMD CDNA wavefront-களில் 64 threads இருப்பதால், அமைதியான தர்க்கப் பிழைகள் (logical bugs) அடிக்கடி ஏற்படுகின்றன.

• Triton ஏன் நான்கு பிரச்சினைகளையும் தவிர்க்கிறது: OpenAI Triton-இல், tl.sum(), tl.max() போன்ற செயல்பாடுகள் நிலைபெயர்த்த 32-thread warp-களுக்குப் பதிலாக முழு logical tensor பரிமாணங்கள் மீதும் இயங்குகின்றன. Compilation-இன்போது, Triton-இன் NVPTX backend, NVIDIA GPU-களுக்குத் தானாகவே 32-thread shuffle trees-ஐ வெளியிடுகிறது; அதேபோல் அதன் AMDGPU backend, AMD Instinct GPU-களுக்கு 64-thread wavefront reductions-ஐ வெளியிடுகிறது. இதனால் warp-size பிழைகள் வடிவமைப்பிலேயே நீக்கப்படுகின்றன.

PyTorch bytecode capture மற்றும் Triton code generation முதல் CUDA எதிர் AMD wavefront இயக்கவியல் வரை, நவீன GPU kernels-இன் முழு வாழ்க்கைச் சுழற்சியையும் இப்போது ஆராய்ந்துள்ளோம்.