Hopper Support in Triton��Jun Yang on behalf of the great team efforts

1

Acknowledgement

• Great collaboration between NVIDIA and OpenAI

​

2

Agenda

3

Recap

Triton makes a trade-off between development efficiency and performance

• More buttons for the geeks
• Users precisely control the behavior of each thread
• More general in different domains
• Longer learning curve & higher development costs
• Best ultimate performance
• Lower lower-bound and higher higher-bound

• ‘Point and shoot’ for ‘primary’ users
• Users describe the computation of each block
• Much less buttons / knobs, easier to get started
• The Triton compiler takes over the optimizations within block level
• Trade-off
• Not suitable for hi-end users with stronger requirements
• Much less buttons/knobs
• Only work for specific domains
• Higher lower-bound but lower higher-bound
• May get more valuable along with the evolution of hardware architecture

​

CUDA

Triton

4

Recap

Triton’s solution on Ampere for the ‘point and shoot’ product positioning

Within 10 lines

5

@triton.jit

def matmul_kernel(

    a_ptr, b_ptr, c_ptr,

    stride_am, stride_ak,

    stride_bk, stride_bn,

    stride_cm, stride_cn,

    M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, # M=N=K=1024

    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, #BLOCK_M=BLOCK_N=BLOCK_K=32

):

    offs_m = tl.arange(0, BLOCK_M)

    offs_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

    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    for k in range(0, K, BLOCK_K):

        a = tl.load(a_ptrs)

        b = tl.load(b_ptrs)

        accumulator += tl.dot(a, b)

        a_ptrs += BLOCK_K * stride_ak

        b_ptrs += BLOCK_K * stride_bk

�    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn

    tl.store(c_ptrs, accumulator)

​

Data Load/Store operations)

Dot operations

Recap

6

module {

  func public @matmul_kernel_0d1d2d3d4c5d6c7d8c(

%arg0: !tt.ptr<f16> {tt.divisibility = 16 : i32}, ...

    %c32_i32 = arith.constant 32 : i32

    %cst = arith.constant dense<32> : tensor<32x32xi32>

    %c32 = arith.constant 32 : index

    %c1024 = arith.constant 1024 : index

    %c0 = arith.constant 0 : index

    %cst_0 = arith.constant dense<0.000000e+00> : tensor<32x32xf32>

    %0 = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<32xi32>

    %1 = tt.expand_dims %0 {axis = 1 : i32} : (tensor<32xi32>) -> tensor<32x1xi32>

    %2 = tt.splat %arg3 : (i32) -> tensor<32x1xi32>

    %3 = arith.muli %1, %2 : tensor<32x1xi32>

...

%11 = arith.muli %1, %10 : tensor<32x1xi32>

    %12 = tt.splat %arg1 : (!tt.ptr<f16>) -> tensor<32x1x!tt.ptr<f16>>

    %13 = tt.addptr %12, %11 : tensor<32x1x!tt.ptr<f16>>, tensor<32x1xi32>

    %14 = tt.broadcast %13 : (tensor<32x1x!tt.ptr<f16>>) -> tensor<32x32x!tt.ptr<f16>>

    %15 = tt.addptr %14, %8 : tensor<32x32x!tt.ptr<f16>>, tensor<32x32xi32>

    %16:3 = scf.for %arg6 = %c0 to %c1024 step %c32 iter_args(%arg7 = %cst_0, %arg8 = %9, %arg9 = %15) -> (...)

      %24 = tt.load %arg8 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<32x32xf16>

      %25 = tt.load %arg9 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<32x32xf16>

      %26 = tt.dot %24, %25, %arg7 {allowTF32 = true} : tensor<32x32xf16> * tensor<32x32xf16> -> tensor<32x32xf32>

      %27 = tt.addptr %arg8, %cst : tensor<32x32x!tt.ptr<f16>>, tensor<32x32xi32>

      %28 = arith.muli %arg4, %c32_i32 : i32

      %29 = tt.splat %28 : (i32) -> tensor<32x32xi32>

      %30 = tt.addptr %arg9, %29 : tensor<32x32x!tt.ptr<f16>>, tensor<32x32xi32>

      scf.yield %26, %27, %30 : tensor<32x32xf32>, tensor<32x32x!tt.ptr<f16>>, tensor<32x32x!tt.ptr<f16>>

    }

    %17 = tt.splat %arg5 : (i32) -> tensor<32x1xi32>

    %18 = arith.muli %1, %17 : tensor<32x1xi32>

    %19 = tt.splat %arg2 : (!tt.ptr<f16>) -> tensor<32x1x!tt.ptr<f16>>

    %20 = tt.addptr %19, %18 : tensor<32x1x!tt.ptr<f16>>, tensor<32x1xi32>

    %21 = tt.broadcast %20 : (tensor<32x1x!tt.ptr<f16>>) -> tensor<32x32x!tt.ptr<f16>>

    %22 = tt.addptr %21, %8 : tensor<32x32x!tt.ptr<f16>>, tensor<32x32xi32>

    %23 = arith.truncf %16#0 : tensor<32x32xf32> to tensor<32x32xf16>

    tt.store %22, %23 {cache = 1 : i32, evict = 1 : i32} : tensor<32x32xf16>

    return

  }

}

Load tile from gmem to shared memory, will get decomposed into 4 vectorized LDGSTS instructions:

​

mov.u32 %r45, 16;

cp.async.cg.shared.global [ %r44 + 0 ], [ %rd11 + 0 ], 0x10, %r45;

cp.async.cg.shared.global [ %r46 + 0 ], [ %rd12 + 0 ], 0x10, %r45;

cp.async.cg.shared.global [ %r48 + 0 ], [ %rd13 + 0 ], 0x10, %r45;

cp.async.cg.shared.global [ %r50 + 0 ], [ %rd14 + 0 ], 0x10, %r45;

mma on tile, will get decomposed into 4 MMA instructions with the corresponding ldmatrix instructions(on Ampere)

ldmatrix.sync.aligned.m8n8.x4.shared.b16 { %r236, %r237, %r238, %r239 }, [ %r56 + 0 ];

…

ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 { %r232, %r233, %r234, %r235 }, [ %r61 + 0 ];

…

ldmatrix.sync.aligned.m8n8.x4.shared.b16 { %r138, %r139, %r140, %r141 }, [ %r120 + 0 ];

…

ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 { %r121, %r122, %r123, %r124 }, [ %r125 + 0 ];

mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %f65, %f66, %f67, %f68 },

{ %r236, %r237, %r238, %r239 }, { %r232, %r233 }, { %f65, %f66, %f67, %f68 };

mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %f69, %f70, %f71, %f72 },

{ %r236, %r237, %r238, %r239 }, { %r234, %r235 }, { %f69, %f70, %f71, %f72 };

mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %f65, %f66, %f67, %f68 },

{ %r138, %r139, %r140, %r141 }, { %r121, %r122 }, { %f65, %f66, %f67, %f68 };

mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %f69, %f70, %f71, %f72 },

{ %r138, %r139, %r140, %r141 }, { %r123, %r124 }, { %f69, %f70, %f71, %f72 };

​

Store from register files to gmem, will get decomposed into 1 vectorized store instruction:

​

st.global.v4.b32 [ %rd32 + 0 ], { %r158, %r159, %r160, %r161 };

Recap

7

Motivation

Triton supports Hopper with decent performance and few additional complexity

H100 speedup over A100 (TC=Tensor Core)

8

Motivation

• Challenges
• New hardware features
• TMA (multicast)
• WGMMA
• Distributed shared memory
• …
• Minimize the changes on frontend APIs
• Backward compatibility
• Future extensibility
• New methodologies in writing kernels
• Heuristics on Ampere might be outdated
• Warp specialization challenges many aspects of Triton

• Changes
• Frontend
• New load/store APIs to adapt to TMA
• Programming on CTA -> Programming on CGA
• Op definition changes of tt/ttgpu dialect
• Backend
• triton::NVGPU dialect
• runtime to maintain TMA descriptors
• WGMMA, TMA ..
• Optimizer pass pipeline
• PlanCTA pass to do the tiling from CGA to CTA
• Backward compatibility of new load/store APIs
• Changes on the existing pass pipeline to adapt to new hardware features
• Warp specialization passes

9

Agenda

10

New features of Hopper

Clusters(CGA) and Distributed Shared Memory

Programming model

Memory Hierarchy

11

New API for CGA programming

Hide the complexity of CGA->CTA tiling and distributed shared memory

​

12

New features of Hopper

TMA Multicast and Asynchronous Transaction Barrier

Asynchronous memory copy with TMA on H100 vs. LDGSTS on A100

Asynchronous barrier in A100 vs. asynchronous transaction barrier in H100

13

New load/store APIs for TMA

14

New Load/Store APIs for TMA

Examples

rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)

rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)

ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)

rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)

rk = pid_z * BLOCK_K + tl.arange(0, BLOCK_K)

# pointers

A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak)

B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn)

acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=dot_out_dtype)

for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)):

if EVEN_K:

a = tl.load(A)

b = tl.load(B)

else:

k_remaining = K - k * (BLOCK_K * SPLIT_K)

a = tl.load(A, mask=rk[None, :] < k_remaining, other=0.)

b = tl.load(B, mask=rk[:, None] < k_remaining, other=0.)

acc += tl.dot(a, b, out_dtype=dot_out_dtype)

A += BLOCK_K * SPLIT_K * stride_ak

B += BLOCK_K * SPLIT_K * stride_bk

A = tl.make_block_ptr(A, shape=(M, K),

strides=(stride_am, stride_ak),

offsets=(pid_m * BLOCK_M, pid_z * BLOCK_K),

block_shape=(BLOCK_M, BLOCK_K),

order=(1, 0))

B = tl.make_block_ptr(B, shape=(K, N),

strides=(stride_bk, stride_bn),

offsets=(pid_z * BLOCK_K, pid_n * BLOCK_N),

block_shape=(BLOCK_K, BLOCK_N),

order=(1, 0))

acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=dot_out_dtype)

for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)):

if EVEN_K:

a = tl.load(A)

b = tl.load(B)

else:

a = tl.load(A, boundary_check=(1, ), padding_option="zero")

b = tl.load(B, boundary_check=(0, ), padding_option="zero")

acc += tl.dot(a, b, out_dtype=dot_out_dtype)

A = tl.advance(A, (0, BLOCK_K * SPLIT_K))

B = tl.advance(B, (BLOCK_K * SPLIT_K, 0))

Legacy APIs

New APIs

15

New Load/Store APIs for TMA

Examples

rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)

rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)

ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)

rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)

rk = pid_z * BLOCK_K + tl.arange(0, BLOCK_K)

# pointers

A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak)

B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn)

acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=dot_out_dtype)

for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)):

if EVEN_K:

a = tl.load(A)

b = tl.load(B)

else:

k_remaining = K - k * (BLOCK_K * SPLIT_K)

a = tl.load(A, mask=rk[None, :] < k_remaining, other=0.)

b = tl.load(B, mask=rk[:, None] < k_remaining, other=0.)

acc += tl.dot(a, b, out_dtype=dot_out_dtype)

A += BLOCK_K * SPLIT_K * stride_ak

B += BLOCK_K * SPLIT_K * stride_bk

A = tl.make_block_ptr(A, shape=(M, K),

strides=(stride_am, stride_ak),

offsets=(pid_m * BLOCK_M, pid_z * BLOCK_K),

block_shape=(BLOCK_M, BLOCK_K),

order=(1, 0))

B = tl.make_block_ptr(B, shape=(K, N),

strides=(stride_bk, stride_bn),

offsets=(pid_z * BLOCK_K, pid_n * BLOCK_N),

block_shape=(BLOCK_K, BLOCK_N),

order=(1, 0))

acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=dot_out_dtype)

for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)):

a = tl.load(A, boundary_check=(1, ), padding_option="zero")

b = tl.load(B, boundary_check=(0, ), padding_option="zero")

acc += tl.dot(a, b, out_dtype=dot_out_dtype)

A = tl.advance(A, (0, BLOCK_K * SPLIT_K))

B = tl.advance(B, (BLOCK_K * SPLIT_K, 0))

Legacy APIs

New APIs could be even simpler. This brings no performance drop on Hopper, and minor performance drop on pre-Hopper

16

New API for Automatic Warp Specialization(Preview)

Better overlapping between data movement and math instructions

Tile2

Prelogue

Tile3

Prelogue

Tile4

Prelogue

• The peak perf of Hopper is up to 6x faster than Ampere, but the bandwidth has not increased so much.
• Overlapping between data movement and math instructions as much as possible is vital for good performance.

17

How to use

https://github.com/openai/triton

18

Agenda

19

Triton work flow

Frontend

20

Triton work flow

New added dialects and attributes

• CTA -> CGA (BlockedEncoding, MMAEncoding, SharedEncoding)
• CTAsPerCGA, CTASplitNum, CTAOrder
• WGMMA (MMAEncoding, SharedEncoding)
• versionMajor, hasLeadingOffset

21

Triton work flow

Overview of pass pipelines

coalesce

plan_cta

accelerate_matmul

optimize_epilogue

pipeline

materialize_load_store

prefetch

remove_layout_conversions

reorder_instructions

fence_insertion

warp_specialize

CSE, DCE …

enable_warp_specialization=False

enable_warp_specialization=True

A new pass for CGA->CTA tiling & optimization to avoid unnecessary dsmem traffic. This also implicitly make optimization strategies for TMA multicast

Assign MmaEncodingAttr for dots. We added MMAv3 support to this pass.

A new optional pass to bypass shared memory and store to global memory from MMAEncoding directly

Pipeline loops to overlapping between loads and dots, we make this pass pipeline WGMMA dots and TMA loads with async semantics.

A couple of new passes for auto warp specialization.

A new pass to materialize remaining loads/stores after pipeline pass with proper instructions.

Remove redundant layout conversions, We make it adapt for GMMA in case A/B from shared memory.

A new pass to insert proper fence between operations who use different memory proxies, LoadOp and DotAsyncOp .e.g.

22

Triton work flow

Hide the complexity of TMA multicast & distributed shared memory

​

• PlanCTA Pass (CGA->CTA tiling)
• Select key ops and assign layout
• Propagate to other ops, minimizing the number of layout conversion across CTAs

…

…

…

key op (dot, reduce)

23

Triton work flow

Better overlapping between data movement and math instructions

decompose

pipeline

Tile2

Prelogue

Tile3

Prelogue

Tile4

Prelogue

enable_warp_specialization=False

enable_warp_specialization=True

24

Triton work flow

Backend

​

TritonGPU/TritonNvidiaGpu Dialect

LLVM/NVVM Dialect

with inline PTX

LLVM/NVVM IR

with inline PTX

PTX

• A temporary solution since NVVM is not ready to support Hopper features
• May be merged into NVVM dialect/NVVM IR ideally in the future

TritonGPU/TritonNvidiaGpu Dialect

LLVM/NVVM Dialect

LLVM/NVVM IR

PTX

25

Agenda

26

Performance of float16 GEMM

0

0.5

1

1.5

2

Speedup

1

1.41

2

Without SplitK shapes

30% perf gap

1.49

26% perf gap

• A100 SXM 80G, H100 SXM 80G
• SplitK GEMM can be implemented with the Triton front end APIs, but it is not our focus for now.

1

1.06

1.84

1.12

With SplitK shapes

27

Performance of Triton on Hopper

Vector Add

Fused Softmax

https://github.com/openai/triton/tree/main/python/tutorials

28

Agenda

29

Plans for the future

• Performance optimization of the warp specialized kernel.
• Turn on the new features on Hopper by default in Triton.
• Optimizing other kernels, e.g., flash-attention.
• …

30

Thank you

31

• Triton pre-Hopper kernels can run on Hopper by re-compilation, no changes necessary.

​

Motivation

32