(example_overlapping-accumulator)= # Overlapping Accumulator Block-Scaled Matrix Multiplication This example demonstrates overlapping accumulator storage in TMEM for 2CTA block scale MMA by comparing double buffered overlapping BLOCK_N=256 accumulators against a single buffered BLOCK_N=256 accumulator. ```python import argparse import pytest import torch import triton import triton.experimental.gluon as gluon import triton.experimental.gluon.language as gl from triton.tools.mxfp import MXFP4Tensor, MXScaleTensor from triton.experimental.gluon.nvidia.blackwell import TensorDescriptor from triton.experimental.gluon.language.nvidia.blackwell import ( TensorMemoryLayout, TensorMemoryScalesLayout, allocate_tensor_memory, tensor_memory_descriptor, tcgen05_copy, tcgen05_commit, tcgen05_mma_barrier_count, tcgen05_mma_scaled, mbarrier, tma, ) def is_blackwell(): if not torch.cuda.is_available(): return False target = triton.runtime.driver.active.get_current_target() return target.backend == "cuda" and torch.cuda.get_device_capability()[0] == 10 # --------------------------------------------------------------------------- # Fixed Configuration # --------------------------------------------------------------------------- FORMATS = [("mxfp8", "mxfp8"), ("mxfp8", "mxfp4"), ("mxfp4", "mxfp4"), ("nvfp4", "nvfp4")] BLOCK_M = 256 BLOCK_N = 256 EPILOGUE_BLOCK_N = 32 NUM_BUFFERS = 5 NUM_CTAS = 2 CGA_LAYOUT = ((1, 0), ) OUT_DTYPE = torch.float16 TMEM_COLS = gl.constexpr(512) # --------------------------------------------------------------------------- # Overlapping Accumulator Layout # --------------------------------------------------------------------------- # yapf: disable # # With 512 TMEM columns, double buffered BLOCK_N=256 accumulators fully consume TMEM # leaving no room for the A and B scales required for a block scale MMA workload. # # Gluon example 04 uses a single buffered BLOCK_N=256 accumulator, leaving # some performance on the table. # # This example keeps the double buffered BLOCK_N=256 accumulators by overlapping them # in TMEM by exactly the size of the A and B scales as shown here: # # TMEM: BLOCK_N TMEM_COLS # 0 256 512 # [---------------------------------------------------------------) # left accum [-------------------------------) # right accum [-------------------------------) # A scales [-) # B scales [---) # # The overlap path demonstrates three major enhancements over example 04: # 1) Fully allocate TMEM and then slice it to hold accumulators and scales. # 2) Epilogue drains overlapped columns first so the next MMA can proceed. # 3) New ping-pong style barrier guards the overlapped columns. # # Tile 0 MMA starts immediately and writes the left accumulator: # # left accum [xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx) # start -> MMA # # Tile 0 epilogue drains the overlap first and releases the overlap barrier: # # <--------- drain from high to low # left accum [xxxxxxxxxxxxxxxxxxxxxxxxx------) # drain overlap -> arrive overlap_bar # # Tile 1 MMA waits for Tile 0 epilogue to release the overlap barrier, then writes the right accumulator: # # left accum [xxxxxxxxxxxxxxxxxxxxxxxxx------) # right accum [xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx) # wait overlap_bar -> MMA # # Tile 1 epilogue drains the overlap first and releases the overlap barrier for Tile 2: # # left accum [xxxxxxxxxxxxxxxxxxx------------) # right accum [-----xxxxxxxxxxxxxxxxxxxxxxxxxx) # drain from low to high --------> # drain overlap -> arrive overlap_bar # # The overlap barrier is now released for Tile 2, but Tile 2 MMA must wait for Tile 0 epilogue # to fully drain the left accumulator. The legacy acc_empty / acc_ready barriers are used to # communicate empty / full accumulator semantics between MMA and epilogue partitions. # # yapf: enable class TmemOverlapStoragePlan: def __init__(self, left_accum_offset, left_accum_cols, right_accum_offset, right_accum_cols, a_scale_offset, a_scale_cols, b_scale_offset, b_scale_cols, early_release_subtile): self.left_accum_offset = left_accum_offset self.left_accum_cols = left_accum_cols self.right_accum_offset = right_accum_offset self.right_accum_cols = right_accum_cols self.a_scale_offset = a_scale_offset self.a_scale_cols = a_scale_cols self.b_scale_offset = b_scale_offset self.b_scale_cols = b_scale_cols self.early_release_subtile = early_release_subtile @gluon.constexpr_function def tmem_overlap_storage_plan(BLOCK_K, EPILOGUE_BLOCK_N, VEC_SIZE, USE_OVERLAP: gl.constexpr): # Hard code BLOCK_N to 256 for this example. BLOCK_N: gl.constexpr = 256 # Scale views have K / VEC_SIZE columns. A scales are M-sharded # across CTAs, so each CTA stores one TMEM column per scale-K column. # B scales are duplicated across CTAs by the scale layout; per CTA, # they store one scale-K column group for each 128-column N block. scale_k = BLOCK_K // VEC_SIZE a_scale_cols = scale_k b_scale_cols = scale_k * (BLOCK_N // 128) # The overlap path uses double buffered BLOCK_N=256 accumulators. if USE_OVERLAP: # The overlap is the combined size of the A and B scales # which cannot exceed the size of BLOCK_N. overlap_cols = a_scale_cols + b_scale_cols assert overlap_cols < BLOCK_N # The left accumulator begins at 0 and ends at BLOCK_N. # The right accumulator begins at BLOCK_N minus the overlap. right_accum_offset = BLOCK_N - overlap_cols # Scale views begin after the overlapping accumulators. a_scale_offset = TMEM_COLS - overlap_cols # Identify the zero-based index of the epilogue subtile containing the last overlap column. early_release_subtile = (overlap_cols - 1) // EPILOGUE_BLOCK_N # The comparison path uses a single buffered BLOCK_N=256 accumulator. else: right_accum_offset = 0 a_scale_offset = BLOCK_N early_release_subtile = 0 return TmemOverlapStoragePlan(0, BLOCK_N, right_accum_offset, BLOCK_N, a_scale_offset, a_scale_cols, a_scale_offset + a_scale_cols, b_scale_cols, early_release_subtile) @gluon.constexpr_function def tmem_overlap_storage_plan_from_partition(p): A_ELEM_PER_BYTE: gl.constexpr = 2 if p.a_desc.dtype == gl.uint8 else 1 BLOCK_N: gl.constexpr = p.b_desc.block_shape[0] assert BLOCK_N == 256 BLOCK_K: gl.constexpr = p.a_desc.block_shape[1] * A_ELEM_PER_BYTE EPILOGUE_BLOCK_N: gl.constexpr = p.c_desc.block_shape[1] VEC_SIZE: gl.constexpr = 32 if p.a_scale_desc.dtype == gl.uint8 else 16 return tmem_overlap_storage_plan(BLOCK_K, EPILOGUE_BLOCK_N, VEC_SIZE, p.USE_OVERLAP) def test_tmem_overlap_accum(): plan = tmem_overlap_storage_plan(BLOCK_K=256, EPILOGUE_BLOCK_N=32, VEC_SIZE=16, USE_OVERLAP=True) assert (plan.left_accum_offset, plan.left_accum_cols) == (0, 256) assert (plan.right_accum_offset, plan.right_accum_cols) == (208, 256) assert (plan.a_scale_offset, plan.a_scale_cols) == (464, 16) assert (plan.b_scale_offset, plan.b_scale_cols) == (480, 32) assert plan.early_release_subtile == 1 plan = tmem_overlap_storage_plan(BLOCK_K=128, EPILOGUE_BLOCK_N=32, VEC_SIZE=32, USE_OVERLAP=True) assert (plan.right_accum_offset, plan.right_accum_cols) == (244, 256) assert (plan.a_scale_offset, plan.a_scale_cols) == (500, 4) assert (plan.b_scale_offset, plan.b_scale_cols) == (504, 8) assert plan.early_release_subtile == 0 plan = tmem_overlap_storage_plan(BLOCK_K=256, EPILOGUE_BLOCK_N=32, VEC_SIZE=16, USE_OVERLAP=False) assert (plan.left_accum_offset, plan.left_accum_cols) == (0, 256) assert (plan.right_accum_offset, plan.right_accum_cols) == (0, 256) assert (plan.a_scale_offset, plan.a_scale_cols) == (256, 16) assert (plan.b_scale_offset, plan.b_scale_cols) == (272, 32) # --------------------------------------------------------------------------- # Utilities # --------------------------------------------------------------------------- def random_quantized_tensor(MN, K, format): assert format in ["mxfp4", "mxfp8", "nvfp4"] VEC_SIZE = 16 if format == "nvfp4" else 32 # Generate a random quantized tensor and its scale factors, assuming we are # scaling along the K dimension. base = MXFP4Tensor(size=(MN, K), device="cuda").random() scale = MXScaleTensor(size=(MN, K // VEC_SIZE), device="cuda").random(low=1 / 128, high=2.0) # Compute the dequantized tensor to use for testing. ref = base.to(torch.float32) scale_ref = scale.to(torch.float32) value = ref * scale_ref.repeat_interleave(VEC_SIZE, dim=1) if format == "mxfp8": # For mxfp8, convert the tensor to a regular float8 torch tensor. return ref.to(torch.float8_e4m3fn), scale.data, value elif format == "mxfp4": # For mxfp4, pack the elements along the K dimension. return base.to_packed_tensor(dim=1), scale.data, value else: # For nvfp4, pack the elements along the K dimension, and convert the # scale factors to float8_e4m3fn. return base.to_packed_tensor(dim=1), scale_ref.to(torch.float8_e4m3fn), value def align_to(a, b): # Return next multiple of `b` greater than or equal to `a`. return triton.cdiv(a, b) * b def swizzle_scales_packed_block(scales: torch.Tensor): # When the scale tensor is not an even multiple of [128, 4], we need to pad # the scale tensor so it can use the packed block format. PAD_MN = align_to(scales.shape[0], 128) - scales.shape[0] PAD_K = align_to(scales.shape[1], 4) - scales.shape[1] scales = torch.nn.functional.pad(scales, (0, PAD_K, 0, PAD_MN)) MN, SCALE_K = scales.shape[0], scales.shape[1] REP_MN = MN // 128 REP_K = SCALE_K // 4 scales = scales.reshape(REP_MN, 4, 32, REP_K, 4) scales = scales.permute(0, 3, 2, 1, 4) return scales.contiguous() @gluon.jit def unswizzle_scales_shared_memory(smem, BLOCK_MN: gl.constexpr, BLOCK_K: gl.constexpr, VEC_SIZE: gl.constexpr): smem = smem.reshape((smem.shape[1], smem.shape[2], 32, 4, 4)) smem = smem.permute((0, 3, 2, 1, 4)) return smem.reshape((BLOCK_MN, BLOCK_K // VEC_SIZE)) @gluon.jit def async_mma_scaled_impl(a_smem, b_smem, a_scale_smem, b_scale_smem, acc_tmem, tmem_pool, tmem_plan, mma_bar, use_acc, pred): A_ELEM_PER_BYTE: gl.constexpr = 2 if a_smem.dtype == gl.uint8 else 1 BLOCK_M: gl.constexpr = a_smem.shape[0] BLOCK_N: gl.constexpr = b_smem.shape[0] BLOCK_K: gl.constexpr = a_smem.shape[1] * A_ELEM_PER_BYTE # Recall we use `uint8` to represent fp4 elements. VEC_SIZE: gl.constexpr = 32 if a_scale_smem.dtype == gl.uint8 else 16 a_scale = unswizzle_scales_shared_memory(a_scale_smem, BLOCK_M, BLOCK_K, VEC_SIZE) b_scale = unswizzle_scales_shared_memory(b_scale_smem, BLOCK_N, BLOCK_K, VEC_SIZE) two_ctas: gl.constexpr = acc_tmem.type.layout.two_ctas a_scale_layout: gl.constexpr = TensorMemoryScalesLayout(cga_layout=[[1, 0]] if two_ctas else []) b_scale_layout: gl.constexpr = TensorMemoryScalesLayout(cga_layout=[[0, 0]] if two_ctas else []) a_scale_tmem = tmem_pool.slice(tmem_plan.a_scale_offset, tmem_plan.a_scale_cols)._reinterpret(dtype=a_scale.dtype, shape=a_scale.shape, layout=a_scale_layout) b_scale_tmem = tmem_pool.slice(tmem_plan.b_scale_offset, tmem_plan.b_scale_cols)._reinterpret(dtype=b_scale.dtype, shape=b_scale.shape, layout=b_scale_layout) tcgen05_copy(a_scale, a_scale_tmem) tcgen05_copy(b_scale, b_scale_tmem) a_format: gl.constexpr = "e2m1" if a_smem.dtype == gl.uint8 else "e4m3" b_format: gl.constexpr = "e2m1" if b_smem.dtype == gl.uint8 else "e4m3" tcgen05_mma_scaled(a_smem, b_smem.permute((1, 0)), acc_tmem, a_scale_tmem, b_scale_tmem, a_format, b_format, use_acc=use_acc, pred=pred, multicast=True, mbarriers=[mma_bar]) @gluon.jit def issue_loads(producer, pid_m, pid_n, k, a_desc, b_desc, a_scale_desc, b_scale_desc, a_bufs, b_bufs, a_scale_bufs, b_scale_bufs, bars, pred): A_ELEM_PER_BYTE: gl.constexpr = 2 if a_desc.dtype == gl.uint8 else 1 B_ELEM_PER_BYTE: gl.constexpr = 2 if b_desc.dtype == gl.uint8 else 1 BLOCK_M: gl.constexpr = a_desc.block_shape[0] BLOCK_N: gl.constexpr = b_desc.block_shape[0] BLOCK_K: gl.constexpr = a_desc.block_shape[1] * A_ELEM_PER_BYTE REP_M: gl.constexpr = a_scale_desc.block_shape[1] REP_N: gl.constexpr = b_scale_desc.block_shape[1] A_REP_K: gl.constexpr = a_scale_desc.block_shape[2] B_REP_K: gl.constexpr = b_scale_desc.block_shape[2] off_m = pid_m * BLOCK_M off_n = pid_n * BLOCK_N off_m_a_scale = pid_m * REP_M off_n_b_scale = pid_n * REP_N off_k_a = k // A_ELEM_PER_BYTE off_k_b = k // B_ELEM_PER_BYTE off_k_a_scale = (k // BLOCK_K) * A_REP_K off_k_b_scale = (k // BLOCK_K) * B_REP_K index = producer.index bar = bars.index(index) mbarrier.expect( bar, a_desc.nbytes_per_cta + b_desc.nbytes_per_cta + a_scale_desc.nbytes_per_cta + b_scale_desc.nbytes_per_cta, pred) tma.async_load(a_desc, [off_m, off_k_a], bar, a_bufs.index(index), pred, multicast=True) tma.async_load(b_desc, [off_n, off_k_b], bar, b_bufs.index(index), pred, multicast=True) tma.async_load(a_scale_desc, [0, off_m_a_scale, off_k_a_scale, 0, 0], bar, a_scale_bufs.index(index), pred, multicast=True) tma.async_load(b_scale_desc, [0, off_n_b_scale, off_k_b_scale, 0, 0], bar, b_scale_bufs.index(index), pred, multicast=True) return producer.next(pred) @gluon.jit def issue_mma(consumer, c_bars, a_bufs, b_bufs, a_scale_bufs, b_scale_bufs, producer, p_bars, acc_tmem, tmem_pool, tmem_plan, use_acc, pred): c_index = consumer.index mbarrier.wait(c_bars.index(c_index), consumer.phase, pred) async_mma_scaled_impl(a_bufs.index(c_index), b_bufs.index(c_index), a_scale_bufs.index(c_index), b_scale_bufs.index(c_index), acc_tmem, tmem_pool, tmem_plan, p_bars.index(producer.index), use_acc, pred) return consumer.next(pred), producer.next(pred) @gluon.aggregate class Counter: index: gl.tensor phase: gl.tensor num_barriers: gl.constexpr @gluon.jit def create(phase, num_barriers: gl.constexpr): return Counter(gl.to_tensor(0), gl.to_tensor(phase), num_barriers) @gluon.must_use_result @gluon.jit def next(self, pred=True): incr = self.index + gl.where(pred, 1, 0) rollover = incr == self.num_barriers index = gl.where(rollover, 0, incr) phase = gl.where(rollover, self.phase ^ 1, self.phase) return Counter(index, phase, self.num_barriers) @gluon.aggregate class SpsTileScheduler: has_work: gl.tensor tile_id: gl.tensor pid_m: gl.tensor pid_n: gl.tensor TILE_M: gl.constexpr TILE_N: gl.constexpr NUM_PID_M: gl.tensor NUM_PID_N: gl.tensor @gluon.jit def initialize(TILE_M: gl.constexpr, TILE_N: gl.constexpr, NUM_PID_M, NUM_PID_N): tile_id = gl.program_id(axis=0) has_work = tile_id < NUM_PID_M * NUM_PID_N pid_m = gl.to_tensor(0) pid_n = gl.to_tensor(0) if has_work: pid_m = tile_id // NUM_PID_N pid_n = tile_id % NUM_PID_N return SpsTileScheduler(has_work, tile_id, pid_m, pid_n, TILE_M, TILE_N, NUM_PID_M, NUM_PID_N) @gluon.jit def get_offsets(self): return self.pid_m * self.TILE_M, self.pid_n * self.TILE_N @gluon.jit def step(self): next_tile_id = self.tile_id + gl.num_programs(axis=0) has_work = next_tile_id < self.NUM_PID_M * self.NUM_PID_N pid_m = gl.to_tensor(0) pid_n = gl.to_tensor(0) if has_work: pid_m = next_tile_id // self.NUM_PID_N pid_n = next_tile_id % self.NUM_PID_N return SpsTileScheduler(has_work, next_tile_id, pid_m, pid_n, self.TILE_M, self.TILE_N, self.NUM_PID_M, self.NUM_PID_N) # --------------------------------------------------------------------------- # Partitions # --------------------------------------------------------------------------- @gluon.aggregate class PartitionArgs: a_desc: tma.tensor_descriptor b_desc: tma.tensor_descriptor c_desc: tma.tensor_descriptor a_scale_desc: tma.tensor_descriptor b_scale_desc: tma.tensor_descriptor a_bufs: gl.shared_memory_descriptor b_bufs: gl.shared_memory_descriptor a_scale_bufs: gl.shared_memory_descriptor b_scale_bufs: gl.shared_memory_descriptor load_empty_bars: gl.shared_memory_descriptor load_ready_bars: gl.shared_memory_descriptor acc_bufs: tensor_memory_descriptor acc_empty_bars: gl.shared_memory_descriptor acc_ready_bars: gl.shared_memory_descriptor overlap_bar: gl.shared_memory_descriptor NUM_PID_M: gl.tensor NUM_PID_N: gl.tensor USE_OVERLAP: gl.constexpr @gluon.jit def get_sps_scheduler(self): return SpsTileScheduler.initialize( self.a_desc.block_shape[0], self.b_desc.block_shape[0], self.NUM_PID_M, self.NUM_PID_N, ) @gluon.jit def mma_scaled_load_partition(p): A_ELEM_PER_BYTE: gl.constexpr = 2 if p.a_desc.dtype == gl.uint8 else 1 BLOCK_K: gl.constexpr = p.a_desc.block_shape[1] * A_ELEM_PER_BYTE K = p.a_desc.shape[1] * A_ELEM_PER_BYTE state = Counter.create(1, p.load_empty_bars.shape[0]) scheduler = p.get_sps_scheduler() while scheduler.has_work: for k in range(0, K, BLOCK_K): mbarrier.wait(p.load_empty_bars.index(state.index), state.phase) state = issue_loads(state, scheduler.pid_m, scheduler.pid_n, k, p.a_desc, p.b_desc, p.a_scale_desc, p.b_scale_desc, p.a_bufs, p.b_bufs, p.a_scale_bufs, p.b_scale_bufs, p.load_ready_bars, pred=True) scheduler = scheduler.step() @gluon.jit def mma_scaled_mma_partition(p): A_ELEM_PER_BYTE: gl.constexpr = 2 if p.a_desc.dtype == gl.uint8 else 1 BLOCK_K: gl.constexpr = p.a_desc.block_shape[1] * A_ELEM_PER_BYTE K = p.a_desc.shape[1] * A_ELEM_PER_BYTE tmem_plan: gl.constexpr = tmem_overlap_storage_plan_from_partition(p) load_state = Counter.create(0, p.load_empty_bars.shape[0]) acc_state = Counter.create(1, p.acc_empty_bars.shape[0]) scheduler = p.get_sps_scheduler() i = 0 while scheduler.has_work: # The overlap barrier only protects the shared columns. Before reusing an # accumulator window, wait until its epilogue has drained the entire window. # For example, Tile 2 cannot reuse the left window until Tile 0 releases it. mbarrier.wait(p.acc_empty_bars.index(acc_state.index), acc_state.phase) # MMA must wait for the preceding epilogue to drain the overlap. if p.USE_OVERLAP: mbarrier.wait(p.overlap_bar.index(0), (i + 1) % 2) # When overlap is disabled, always use the left (only) accumulator. # When overlap is enabled, even tiles use the left and odd tiles use the right accumulator. if p.USE_OVERLAP and i % 2 != 0: acc_tmem = p.acc_bufs.slice(tmem_plan.right_accum_offset, tmem_plan.right_accum_cols) else: acc_tmem = p.acc_bufs.slice(tmem_plan.left_accum_offset, tmem_plan.left_accum_cols) use_acc = False for k in range(0, K, BLOCK_K): _, load_state = issue_mma(load_state, p.load_ready_bars, p.a_bufs, p.b_bufs, p.a_scale_bufs, p.b_scale_bufs, load_state, p.load_empty_bars, acc_tmem, p.acc_bufs, tmem_plan, use_acc, pred=True) use_acc = True tcgen05_commit(p.acc_ready_bars.index(acc_state.index)) acc_state = acc_state.next() scheduler = scheduler.step() i += 1 @gluon.jit def mma_scaled_epilogue_partition(p): tile_m: gl.constexpr = p.c_desc.block_shape[0] BLOCK_N: gl.constexpr = p.b_desc.block_shape[0] EPILOGUE_BLOCK_N: gl.constexpr = p.c_desc.block_shape[1] tmem_plan: gl.constexpr = tmem_overlap_storage_plan_from_partition(p) subtile_factor: gl.constexpr = BLOCK_N // EPILOGUE_BLOCK_N subtile_stages: gl.constexpr = 1 if subtile_factor == 1 else 2 acc_state = Counter.create(0, p.acc_empty_bars.shape[0]) acc_smems = gl.allocate_shared_memory(p.c_desc.dtype, [subtile_stages, tile_m, EPILOGUE_BLOCK_N], p.c_desc.layout) sub_acc_state = Counter.create(0, subtile_stages) scheduler = p.get_sps_scheduler() i = 0 while scheduler.has_work: off_m, off_n = scheduler.get_offsets() mbarrier.wait(p.acc_ready_bars.index(acc_state.index), acc_state.phase) # When overlap is disabled, always use the left (only) accumulator. # When overlap is enabled, even tiles use the left and odd tiles use the right accumulator. if p.USE_OVERLAP and i % 2 != 0: acc_tmem = p.acc_bufs.slice(tmem_plan.right_accum_offset, tmem_plan.right_accum_cols) else: acc_tmem = p.acc_bufs.slice(tmem_plan.left_accum_offset, tmem_plan.left_accum_cols) for s in gl.static_range(subtile_factor): # When overlap is disabled, always drain subtiles from low-to-high. # When overlap is enabled, drain the overlap first: # Even tiles use the left accumulator and drain high-to-low. # Odd tiles use the right accumulator and drain low-to-high. if p.USE_OVERLAP and i % 2 == 0: acc_sub = acc_tmem.slice(EPILOGUE_BLOCK_N * (subtile_factor - 1 - s), EPILOGUE_BLOCK_N) store_n = off_n + EPILOGUE_BLOCK_N * (subtile_factor - 1 - s) else: acc_sub = acc_tmem.slice(EPILOGUE_BLOCK_N * s, EPILOGUE_BLOCK_N) store_n = off_n + EPILOGUE_BLOCK_N * s acc_smem = acc_smems.index(sub_acc_state.index) acc = acc_sub.load().to(p.c_desc.dtype) # Signal the barrier once the epilogue drains the overlap so the next MMA can proceed. if p.USE_OVERLAP and s == tmem_plan.early_release_subtile: mbarrier.arrive(p.overlap_bar.index(0), count=1) tma.store_wait(pendings=subtile_stages - 1) acc_smem.store(acc) tma.async_copy_shared_to_global(p.c_desc, [off_m, store_n], acc_smem) sub_acc_state = sub_acc_state.next() mbarrier.arrive(p.acc_empty_bars.index(acc_state.index), count=1) acc_state = acc_state.next() scheduler = scheduler.step() i += 1 tma.store_wait(0) # --------------------------------------------------------------------------- # Kernel # --------------------------------------------------------------------------- @gluon.jit def mma_scaled_warp_specialized_kernel(a_desc, b_desc, c_desc, a_scale_desc, b_scale_desc, M, N, K, A_ELEM_PER_BYTE, num_buffers: gl.constexpr, BLOCK_M: gl.constexpr, BLOCK_N: gl.constexpr, BLOCK_K: gl.constexpr, EPILOGUE_BLOCK_N: gl.constexpr, CGA_LAYOUT: gl.constexpr, USE_OVERLAP: gl.constexpr): NUM_CTAS: gl.constexpr = gl.num_ctas() TWO_CTAS: gl.constexpr = NUM_CTAS > 1 BLOCK_M_PER_CTA: gl.constexpr = BLOCK_M // NUM_CTAS gl.static_assert(BLOCK_M_PER_CTA == 64 or BLOCK_M_PER_CTA == 128) num_acc_buffers: gl.constexpr = 2 if USE_OVERLAP else 1 a_bufs = gl.allocate_shared_memory(a_desc.dtype, [num_buffers] + a_desc.block_shape, a_desc.layout) b_bufs = gl.allocate_shared_memory(b_desc.dtype, [num_buffers] + b_desc.block_shape, b_desc.layout) a_scale_bufs = gl.allocate_shared_memory(a_scale_desc.dtype, [num_buffers] + a_scale_desc.block_shape, a_scale_desc.layout) b_scale_bufs = gl.allocate_shared_memory(b_scale_desc.dtype, [num_buffers] + b_scale_desc.block_shape, b_scale_desc.layout) tmem_layout: gl.constexpr = TensorMemoryLayout([BLOCK_M_PER_CTA, BLOCK_N], col_stride=1, cga_layout=CGA_LAYOUT, two_ctas=TWO_CTAS) # Allocate all of TMEM and then slice it (later) to accommodate accumulators and scales. acc_bufs = allocate_tensor_memory(gl.float32, [BLOCK_M, TMEM_COLS], tmem_layout) mma_barrier_count: gl.constexpr = tcgen05_mma_barrier_count( [a_bufs.index(0), b_bufs.index(0), a_scale_bufs.index(0), b_scale_bufs.index(0)], multicast=True, two_ctas=TWO_CTAS) load_empty_bars = mbarrier.allocate_mbarrier(batch=num_buffers) load_ready_bars = mbarrier.allocate_mbarrier(batch=num_buffers, two_ctas=TWO_CTAS) for i in gl.static_range(num_buffers): mbarrier.init(load_empty_bars.index(i), count=mma_barrier_count) mbarrier.init(load_ready_bars.index(i), count=1) acc_empty_bars = mbarrier.allocate_mbarrier(batch=num_acc_buffers, two_ctas=TWO_CTAS) acc_ready_bars = mbarrier.allocate_mbarrier(batch=num_acc_buffers) for i in gl.static_range(num_acc_buffers): mbarrier.init(acc_empty_bars.index(i), count=1) mbarrier.init(acc_ready_bars.index(i), count=1) overlap_bar = mbarrier.allocate_mbarrier(batch=1, two_ctas=TWO_CTAS) mbarrier.init(overlap_bar.index(0), count=1) p = PartitionArgs(a_desc, b_desc, c_desc, a_scale_desc, b_scale_desc, a_bufs, b_bufs, a_scale_bufs, b_scale_bufs, load_empty_bars, load_ready_bars, acc_bufs, acc_empty_bars, acc_ready_bars, overlap_bar, gl.cdiv(M, BLOCK_M), gl.cdiv(N, BLOCK_N), USE_OVERLAP) gl.warp_specialize([ (mma_scaled_epilogue_partition, (p, )), (mma_scaled_mma_partition, (p, )), (mma_scaled_load_partition, (p, )), ], [1, 1], [24, 24]) # --------------------------------------------------------------------------- # Wrapper # --------------------------------------------------------------------------- def make_descriptors(A, B, A_scale, B_scale, M, N): """Create TMA descriptors for the fixed example configuration.""" a_is_fp4 = A.dtype == torch.uint8 b_is_fp4 = B.dtype == torch.uint8 mixed_prec = a_is_fp4 != b_is_fp4 a_elem_per_byte = 2 if a_is_fp4 else 1 b_elem_per_byte = 2 if b_is_fp4 else 1 BLOCK_K = 128 if torch.float8_e4m3fn in [A.dtype, B.dtype] else 256 a_block = [BLOCK_M, BLOCK_K // a_elem_per_byte] b_block = [BLOCK_N, BLOCK_K // b_elem_per_byte] c_block = [BLOCK_M, EPILOGUE_BLOCK_N] cga = tuple(tuple(x) for x in CGA_LAYOUT) a_layout = gl.NVMMASharedLayout.get_default_for(a_block, gl.uint8 if a_is_fp4 else gl.float8e4nv, fp4_padded=(mixed_prec and a_is_fp4), cga_layout=cga) b_layout = gl.NVMMASharedLayout.get_default_for(b_block, gl.uint8 if b_is_fp4 else gl.float8e4nv, fp4_padded=(mixed_prec and b_is_fp4), cga_layout=cga) a_desc = TensorDescriptor.from_tensor(A, a_block, a_layout) b_desc = TensorDescriptor.from_tensor(B, b_block, b_layout) C = torch.empty(M, N, device="cuda", dtype=OUT_DTYPE) C_dtype = getattr(gl, str(OUT_DTYPE).split('.')[1]) c_layout = gl.NVMMASharedLayout.get_default_for(c_block, C_dtype, cga_layout=cga) c_desc = TensorDescriptor.from_tensor(C, c_block, c_layout) is_nvfp4 = A_scale.dtype == torch.float8_e4m3fn vec_size = 16 if is_nvfp4 else 32 rep_m = BLOCK_M // 128 rep_n = BLOCK_N // 128 rep_k = BLOCK_K // (vec_size * 4) a_scale_block = [1, rep_m, rep_k, 2, 256] b_scale_block = [1, rep_n, rep_k, 2, 256] cga_a_scale = [[0, 1, 0, 0, 0]] cga_b_scale = [[0, 0, 0, 0, 0]] a_scale_layout = gl.NVMMASharedLayout(swizzle_byte_width=0, element_bitwidth=8, rank=5, cga_layout=cga_a_scale) b_scale_layout = gl.NVMMASharedLayout(swizzle_byte_width=0, element_bitwidth=8, rank=5, cga_layout=cga_b_scale) A_scale_5d = A_scale.reshape(1, A_scale.shape[0], A_scale.shape[1], 2, 256) B_scale_5d = B_scale.reshape(1, B_scale.shape[0], B_scale.shape[1], 2, 256) a_scale_desc = TensorDescriptor.from_tensor(A_scale_5d, a_scale_block, a_scale_layout) b_scale_desc = TensorDescriptor.from_tensor(B_scale_5d, b_scale_block, b_scale_layout) return a_desc, b_desc, c_desc, a_scale_desc, b_scale_desc, BLOCK_K def mma_scaled_warp_specialized(A, B, A_scale, B_scale, use_overlap=True): """Warp-specialized block-scale MMA with optional overlapping accumulators.""" M, N = A.shape[0], B.shape[0] A_ELEM_PER_BYTE = 2 if A.dtype == torch.uint8 else 1 B_ELEM_PER_BYTE = 2 if B.dtype == torch.uint8 else 1 K = A.shape[1] * A_ELEM_PER_BYTE assert K == B.shape[1] * B_ELEM_PER_BYTE A_desc, B_desc, C_desc, A_scale_desc, B_scale_desc, BLOCK_K = make_descriptors(A, B, A_scale, B_scale, M, N) num_pid = triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N) sm_count = torch.cuda.get_device_properties(A.device).multi_processor_count grid = (max(1, min(num_pid, sm_count // NUM_CTAS)), ) mma_scaled_warp_specialized_kernel[grid]( A_desc, B_desc, C_desc, A_scale_desc, B_scale_desc, M, N, K, A_ELEM_PER_BYTE, NUM_BUFFERS, BLOCK_M, BLOCK_N, BLOCK_K, EPILOGUE_BLOCK_N, CGA_LAYOUT, use_overlap, num_ctas=NUM_CTAS, ) return C_desc.base def make_problem(M, N, K, a_format, b_format): A, A_scale, A_ref = random_quantized_tensor(M, K, a_format) B, B_scale, B_ref = random_quantized_tensor(N, K, b_format) A_scale = swizzle_scales_packed_block(A_scale) B_scale = swizzle_scales_packed_block(B_scale) return A, B, A_scale, B_scale, A_ref @ B_ref.T def benchmark(M, N, K, a_format, b_format, use_overlap): A, B, A_scale, B_scale, _ = make_problem(M, N, K, a_format, b_format) def run(): return mma_scaled_warp_specialized(A, B, A_scale, B_scale, use_overlap=use_overlap) ms = triton.testing.do_bench_cudagraph(run) return 2.0 * M * N * K * 1.0e-9 / ms @pytest.mark.parametrize("a_format, b_format", FORMATS) @pytest.mark.parametrize("use_overlap", [False, True]) @pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell") def test_mma_scaled_overlap_accumulator(a_format, b_format, use_overlap): torch.manual_seed(0) A, B, A_scale, B_scale, C_ref = make_problem(8192, 8192, 8192, a_format, b_format) C = mma_scaled_warp_specialized(A, B, A_scale, B_scale, use_overlap=use_overlap) torch.testing.assert_close(C_ref, C.to(torch.float32), atol=2e-3, rtol=1e-3) def main(): parser = argparse.ArgumentParser() parser.add_argument("--M", type=int, default=8192) parser.add_argument("--N", type=int, default=8192) parser.add_argument("--K", type=int, nargs="+", default=[512, 1024, 2048, 4096, 8192]) args = parser.parse_args() print("format K off on speedup") for a_format, b_format in FORMATS: label = f"{a_format}-{b_format}" for K in args.K: torch.manual_seed(0) off = benchmark(args.M, args.N, K, a_format, b_format, use_overlap=False) torch.manual_seed(0) on = benchmark(args.M, args.N, K, a_format, b_format, use_overlap=True) speedup = (on / off - 1.0) * 100.0 print(f"{label:14s} {K:7d} {off:9.1f} {on:9.1f} {speedup:8.2f}%") print() if __name__ == "__main__": main() ```