#
The execution order of thread blocks in a CUDA kernel affects what data is resident in the L2 cache at any given moment, which in turn influences memory access efficiency and overall kernel performance, especially when memory access becomes the bottleneck of the kernel. The execution order can be controlled using thread block swizzle algorithms, which determine the mapping from program IDs to tile locations. Different thread block swizzle algorithms can have different data access and cache utilization patterns, resulting in CUDA kernel performance variations.
In this blog post, I would like to quickly introduce several common thread block swizzle algorithms, simulate their L2 cache efficiency using a simple model, benchmark the performances of GEMM kernels using different thread block swizzle algorithms, and analyze the results.
#
Unlike the brain-twisting shared memory swizzle, thread block swizzle is relatively straightforward to understand. We could like to use a $8 \times 8$ GEMM tiled grid for the output matrix for illustration purposes. The implementations of different thread block swizzle algorithms will be presented in the benchmark script later.
The swizzled index of a $8 \times 8$ grid in a row-major order would be:
$$
\begin{array}{|c|c|c|c|c|c|c|c|}
\hline
00 & 01 & 02 & 03 & 04 & 05 & 06 & 07 \ \hline
08 & 09 & 10 & 11 & 12 & 13 & 14 & 15 \ \hline
16 & 17 & 18 & 19 & 20 & 21 & 22 & 23 \ \hline
24 & 25 & 26 & 27 & 28 & 29 & 30 & 31 \ \hline
32 & 33 & 34 & 35 & 36 & 37 & 38 & 39 \ \hline
40 & 41 & 42 & 43 & 44 & 45 & 46 & 47 \ \hline
48 & 49 & 50 & 51 & 52 & 53 & 54 & 55 \ \hline
56 & 57 & 58 & 59 & 60 & 61 & 62 & 63 \ \hline
\end{array}
$$
The swizzled index of a $8 \times 8$ grid using a grouped 2D panel swizzle with a group size of 4 would be:
$$
\begin{array}{|c|c|c|c|c|c|c|c|}
\hline
00 & 04 & 08 & 12 & 16 & 20 & 24 & 28 \ \hline
01 & 05 & 09 & 13 & 17 & 21 & 25 & 29 \ \hline
02 & 06 & 10 & 14 & 18 & 22 & 26 & 30 \ \hline
03 & 07 & 11 & 15 & 19 & 23 & 27 & 31 \ \hline
32 & 36 & 40 & 44 & 48 & 52 & 56 & 60 \ \hline
33 & 37 & 41 & 45 & 49 & 53 & 57 & 61 \ \hline
34 & 38 & 42 & 46 & 50 & 54 & 58 & 62 \ \hline
35 & 39 & 43 & 47 & 51 & 55 & 59 & 63 \ \hline
\end{array}
$$
This is the most common thread block swizzle algorithm used for GEMM in practice, apart from the linear swizzle, which is usually default.
The swizzled index of a $8 \times 8$ grid using a Morton / Z-order swizzle would be:
$$
\begin{array}{|c|c|c|c|c|c|c|c|}
\hline
00 & 02 & 08 & 10 & 32 & 34 & 40 & 42 \ \hline
01 & 03 & 09 & 11 & 33 & 35 & 41 & 43 \ \hline
04 & 06 & 12 & 14 & 36 & 38 & 44 & 46 \ \hline
05 & 07 & 13 & 15 & 37 & 39 & 45 & 47 \ \hline
16 & 18 & 24 & 26 & 48 & 50 & 56 & 58 \ \hline
17 & 19 & 25 & 27 & 49 & 51 & 57 & 59 \ \hline
20 & 22 & 28 & 30 & 52 & 54 & 60 & 62 \ \hline
21 & 23 & 29 & 31 & 53 & 55 & 61 & 63 \ \hline
\end{array}
$$
Note that Morton / Z-order swizzle only has a bijective mapping for grid sizes that are powers of two.
The swizzled index of a $8 \times 8$ grid using a bitwise XOR swizzle would be:
$$
\begin{array}{|c|c|c|c|c|c|c|c|}
\hline
00 & 09 & 18 & 27 & 36 & 45 & 54 & 63 \ \hline
01 & 08 & 19 & 26 & 37 & 44 & 55 & 62 \ \hline
02 & 11 & 16 & 25 & 38 & 47 & 52 & 61 \ \hline
03 & 10 & 17 & 24 & 39 & 46 & 53 & 60 \ \hline
04 & 13 & 22 & 31 & 32 & 41 & 50 & 59 \ \hline
05 & 12 & 23 & 30 & 33 & 40 & 51 & 58 \ \hline
06 & 15 & 20 & 29 & 34 & 43 & 48 & 57 \ \hline
07 & 14 & 21 & 28 & 35 & 42 & 49 & 56 \ \hline
\end{array}
$$
Note that bitwise XOR swizzle only has a bijective mapping for grid sizes that are powers of two.
#
The L2 cache hit rate of different thread block swizzle algorithms can be simulated using a simple offline model. The performances of each swizzle algorithm applied on GEMM kernels can also be benchmarked.
|
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806
|
"""Triton GEMM thread-block swizzle benchmark.Benchmarks a single Triton FP16/FP32 GEMM kernel across severalprogram-id -> output-tile ("swizzle") orderings and compares each one'sreal, measured throughput against an offline L2 cache hit-rate prediction.Background----------A GEMM kernel launches one thread block ("program") per output tile. Theorder in which program ids are mapped to (M, N) tile coordinates does notaffect correctness, but it determines which tiles are resident in the L2cache (and which DRAM rows are open) at any given moment during execution.Reordering that mapping to favor cache reuse is a standard GEMMoptimization (see Triton's own matmul tutorial). This script benchmarksfive such swizzle families using self-contained CPU pid-mapping functionsand a matching offline L2 hit-rate model, so the predicted ranking can bechecked against real GPU throughput for the same problem: - Linear (Row-Major / Col-Major) - Grouped 2D Panel (configurable group size) - Morton / Z-order - Bitwise XORPersistent vs. non-persistent launch (--persistent)----------------------------------------------------The kernel always loops internally over a fixed, statically-strided subsetof tiles (`tile_id = pid, pid + NUM_PROGRAMS, pid + 2*NUM_PROGRAMS, ...`),so the two launch styles are really just two choices for NUM_PROGRAMS: - Default (non-persistent): NUM_PROGRAMS == total_tiles, i.e. one program per output tile, one loop iteration per program. Real concurrency is whatever the GPU's block scheduler decides at runtime; the offline L2 predictor has to *assume* a concurrency figure (measured occupancy). - `--persistent`: NUM_PROGRAMS == min(total_tiles, measured occupancy), i.e. exactly as many programs as can be concurrently resident, each looping over several tiles. The concurrently-active tile set at loop iteration i is then *exactly* `{c + i*NUM_PROGRAMS : c in [0, NUM_PROGRAMS)}` -- no more scheduler-order guesswork, no launch/ramp-up overhead from dispatching new blocks, no partial "tail wave" at the end.Empirically, this determinism does not meaningfully change the measuredrow-major-vs-col-major gap or the predicted/measured correlation: thepredictor's blind spot is that it never models DRAM row-buffer locality orbank-level contention, and that gap is present regardless of how preciselythe concurrent tile set is known.Methodology-----------1. `check_correctness` validates every swizzle mode against torch.matmul on a small, fixed grid before any benchmarking runs.2. `measure_concurrent_capacity` compiles the kernel once (without executing it) to derive real occupancy (`num_sms * blocks_per_sm`) from its register and shared-memory footprint, rather than assuming a value.3. `predict_l2_hit_rate` runs an offline, byte-granular LRU cache simulation to estimate each swizzle's L2 hit rate for the given problem shape and concurrency (`min(total_tiles, measured_capacity)`).4. Each swizzle mode is timed with triton.testing.do_bench, which clears the L2 cache before every timed repetition, so measurements always start cold and are not contaminated by a previous algorithm's residual state.5. A Spearman rank correlation between predicted hit rate and measured TFLOPS is reported as a sanity check for the offline model.Known limitations------------------ Morton and XOR are only bijective (visit every output tile exactly once) when the grid (ceil(M/BLOCK_M) x ceil(N/BLOCK_N)) is square and a power of two; otherwise they are automatically skipped with a warning (see `build_algorithm_list`).- The L2 hit-rate model is a coarse, whole-tile LRU trace. It does not model DRAM row-buffer locality, bank-level queuing/contention, or memory coalescing, so a higher predicted hit rate does not always translate into higher measured throughput -- in particular, for a compute-bound configuration, throughput can be flat across a wide range of predicted hit rates.- With `--persistent`, tile-to-program assignment is static (round-robin by tile_id), not dynamic/atomic-counter based, so there is no runtime load balancing across programs for edge-tile masking imbalance.Example------- python triton_gemm_swizzle_bench.py --M 8192 --N 8192 --K 8192 \\ --block-m 128 --block-n 128 --block-k 32 \\ --group-sizes 1 2 4 8 16 32 64 --dtype fp16 [--persistent]"""import argparsefrom collections import OrderedDictimport mathimport subprocessimport sysfrom typing import Any, Callable, Dict, List, Tupleimport torchimport tritonimport triton.language as tlfrom triton.runtime import driver# --- Swizzle mode ids shared between the Triton kernel and the CPU predictor ---MODE_ROW_MAJOR = 0MODE_COL_MAJOR = 1MODE_PANEL = 2MODE_MORTON = 3MODE_XOR = 4TORCH_DTYPES = {"fp16": torch.float16, "fp32": torch.float32}TL_DTYPES = {"fp16": tl.float16, "fp32": tl.float32}DTYPE_BYTES = {"fp16": 2, "fp32": 4}# --- Hardware auto-detection ---def detect_gpu_hardware() -> Dict[str, Any]: """Queries the local system for SM count and L2 cache size. Inspects PyTorch `_C.cuda.DeviceProperties` attributes with priority on `L2_cache_size`, falling back to `nvidia-smi` for the SM count/GPU name and to fixed defaults if neither source is available. """ fallback_sms = 108 fallback_l2_mb = 50.0 # Standard usable CUDA global memory L2 limit warnings = [] gpu_name = "Unknown / CPU Fallback" detected_sms = None detected_l2_mb = None if torch.cuda.is_available(): props = torch.cuda.get_device_properties(0) gpu_name = props.name detected_sms = getattr(props, "multi_processor_count", None) for attr in [ "L2_cache_size", "l2_cache_size", "l2CacheSize", "l2_cache_size_bytes", ]: if hasattr(props, attr): l2_bytes = getattr(props, attr) if l2_bytes > 0: detected_l2_mb = l2_bytes / (1024 * 1024) break if detected_sms is None or gpu_name == "Unknown / CPU Fallback": try: cmd = [ "nvidia-smi", "--query-gpu=name,multiprocessor_count", "--format=csv,noheader,nounits", ] output = subprocess.check_output(cmd).decode( "utf-8").strip().split(",") if len(output) >= 2: gpu_name = output[0].strip() if detected_sms is None: detected_sms = int(output[1].strip()) except Exception: pass if detected_sms is None: detected_sms = fallback_sms warnings.append( f"Unable to query SM count from platform. Using default fallback:" f" {fallback_sms} SMs.") if detected_l2_mb is None: detected_l2_mb = fallback_l2_mb if gpu_name == "Unknown / CPU Fallback": warnings.append( "Unable to query GPU platform or L2 cache size. Using default global" f" limit: {fallback_l2_mb:.2f} MB.") return { "num_sms": detected_sms, "l2_cache_size_mb": detected_l2_mb, "gpu_name": gpu_name, "warnings": warnings, }# --- CPU reference pid-mapping functions (mirrored by get_swizzled_pid below) ---def linear_row_major(pid: int, grid_m: int, grid_n: int) -> Tuple[int, int]: return pid // grid_n, pid % grid_ndef linear_col_major(pid: int, grid_m: int, grid_n: int) -> Tuple[int, int]: return pid % grid_m, pid // grid_mdef panel_swizzle(pid: int, grid_m: int, grid_n: int, group_size_m: int = 8) -> Tuple[int, int]: num_pid_in_group = group_size_m * grid_n group_id = pid // num_pid_in_group first_pid_m = group_id * group_size_m group_size_m_actual = min(grid_m - first_pid_m, group_size_m) by = (pid % num_pid_in_group) // group_size_m_actual bx = first_pid_m + (pid % group_size_m_actual) return bx, bydef morton_swizzle(pid: int, grid_m: int, grid_n: int) -> Tuple[int, int]: bx, by = 0, 0 for i in range(16): bx |= ((pid >> (2 * i)) & 1) << i by |= ((pid >> (2 * i + 1)) & 1) << i return bx % grid_m, by % grid_ndef xor_swizzle(pid: int, grid_m: int, grid_n: int) -> Tuple[int, int]: linear_x = pid % grid_n linear_y = pid // grid_n bx = (linear_x ^ linear_y) % grid_m by = linear_y % grid_n return bx, by@triton.jitdef get_swizzled_pid(pid, num_pid_m, num_pid_n, SWIZZLE_MODE: tl.constexpr, GROUP_SIZE_M: tl.constexpr): """Maps a linear tile id to (pid_m, pid_n) output-tile coordinates. SWIZZLE_MODE selects the ordering (0=row-major, 1=col-major, 2=grouped 2D panel, 3=Morton/Z-order, 4=bitwise XOR); GROUP_SIZE_M is only used by the grouped-panel mode. Mirrors the CPU reference functions above (linear_row_major, linear_col_major, panel_swizzle, morton_swizzle, xor_swizzle) so the offline L2 predictor and the actual kernel agree on tile visitation order. """ if SWIZZLE_MODE == 0: # Linear row-major pid_m = pid // num_pid_n pid_n = pid % num_pid_n elif SWIZZLE_MODE == 1: # Linear col-major pid_m = pid % num_pid_m pid_n = pid // num_pid_m elif SWIZZLE_MODE == 2: # Grouped 2D panel num_pid_in_group = GROUP_SIZE_M * num_pid_n group_id = pid // num_pid_in_group first_pid_m = group_id * GROUP_SIZE_M group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) pid_n = (pid % num_pid_in_group) // group_size_m pid_m = first_pid_m + (pid % group_size_m) elif SWIZZLE_MODE == 3: # Morton / Z-order pid_m = 0 pid_n = 0 for i in range(16): pid_m |= ((pid >> (2 * i)) & 1) << i pid_n |= ((pid >> (2 * i + 1)) & 1) << i pid_m = pid_m % num_pid_m pid_n = pid_n % num_pid_n else: # Bitwise XOR linear_x = pid % num_pid_n linear_y = pid // num_pid_n pid_m = (linear_x ^ linear_y) % num_pid_m pid_n = linear_y % num_pid_n return pid_m, pid_n@triton.jitdef gemm_swizzle_kernel( 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, SWIZZLE_MODE: tl.constexpr, GROUP_SIZE_M: tl.constexpr, DTYPE: tl.constexpr, NUM_PROGRAMS: tl.constexpr,): """Masked, tiled GEMM (C = A @ B) with a swappable program-id swizzle. Exactly NUM_PROGRAMS programs are launched; each one statically owns tile ids `pid, pid + NUM_PROGRAMS, pid + 2*NUM_PROGRAMS, ...`. When NUM_PROGRAMS == total_tiles this degenerates to one tile per program (non-persistent, one loop iteration each); when NUM_PROGRAMS is capped at the GPU's measured concurrent capacity, each program loops over several tiles (persistent). Accumulates in fp32 regardless of input dtype and casts down to DTYPE (shared by A, B, and C) on store; correctness is identical across SWIZZLE_MODE values, only tile visitation order (and therefore cache/memory behavior) changes. """ pid = tl.program_id(axis=0) num_pid_m = tl.cdiv(M, BLOCK_M) num_pid_n = tl.cdiv(N, BLOCK_N) num_tiles = num_pid_m * num_pid_n for tile_id in range(pid, num_tiles, NUM_PROGRAMS): pid_m, pid_n = get_swizzled_pid(tile_id, num_pid_m, num_pid_n, SWIZZLE_MODE, GROUP_SIZE_M) offs_am = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_bn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) offs_k = tl.arange(0, BLOCK_K) a_ptrs = a_ptr + offs_am[:, None] * stride_am + offs_k[ None, :] * stride_ak b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_bn[ None, :] * stride_bn acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for k in range(0, tl.cdiv(K, BLOCK_K)): k_remaining = K - k * BLOCK_K a_mask = (offs_am[:, None] < M) & (offs_k[None, :] < k_remaining) b_mask = (offs_k[:, None] < k_remaining) & (offs_bn[None, :] < N) a = tl.load(a_ptrs, mask=a_mask, other=0.0) b = tl.load(b_ptrs, mask=b_mask, other=0.0) acc = tl.dot(a, b, acc) a_ptrs += BLOCK_K * stride_ak b_ptrs += BLOCK_K * stride_bk c = acc.to(DTYPE) offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) c_ptrs = c_ptr + offs_cm[:, None] * stride_cm + offs_cn[ None, :] * stride_cn c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) tl.store(c_ptrs, c, mask=c_mask)def measure_concurrent_capacity(device: torch.device, block_m: int, block_n: int, block_k: int, dtype: str = "fp16", num_warps: int = 4, num_stages: int = 3) -> int: """Compiles the kernel once (never executes it) and derives resident blocks/SM from its actual register and shared-memory footprint (same technique as the Triton softmax tutorial), instead of guessing/accepting an occupancy value on the CLI. Only compiles, so tiny dummy tensors are enough regardless of the real M/N/K of the benchmark. Returns the number of concurrently-resident CTAs across the whole GPU (`num_sms * blocks_per_sm`). """ properties = driver.active.utils.get_device_properties(device.index) num_sms = properties["multiprocessor_count"] num_regs = properties["max_num_regs"] size_smem = properties["max_shared_mem"] warp_size = properties["warpSize"] # Not exposed by triton's get_device_properties on all versions; torch has it reliably. max_threads_per_sm = torch.cuda.get_device_properties( device.index).max_threads_per_multi_processor torch_dtype = TORCH_DTYPES[dtype] a = torch.empty((block_m, block_k), device=device, dtype=torch_dtype) b = torch.empty((block_k, block_n), device=device, dtype=torch_dtype) c = torch.empty((block_m, block_n), device=device, dtype=torch_dtype) kernel = gemm_swizzle_kernel.warmup( a, b, c, block_m, block_n, block_k, a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1), BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k, SWIZZLE_MODE=MODE_ROW_MAJOR, GROUP_SIZE_M=1, DTYPE=TL_DTYPES[dtype], NUM_PROGRAMS=1, num_warps=num_warps, num_stages=num_stages, grid=(1, ), ) kernel._init_handles() n_regs = kernel.n_regs kernel_smem = kernel.metadata.shared reg_occupancy = num_regs // (n_regs * warp_size * num_warps) smem_occupancy = (size_smem // kernel_smem) if kernel_smem > 0 else reg_occupancy thread_occupancy = max_threads_per_sm // (num_warps * warp_size) blocks_per_sm = max(1, min(reg_occupancy, smem_occupancy, thread_occupancy)) return num_sms * blocks_per_smdef gemm_swizzle(a: torch.Tensor, b: torch.Tensor, swizzle_mode: int, block_m: int, block_n: int, block_k: int, num_programs: int, group_size_m: int = 8, dtype: str = "fp16") -> torch.Tensor: """Allocates the output tensor and launches exactly `num_programs` programs for one swizzle mode. Returns the (M, N) result tensor. """ M, K = a.shape K2, N = b.shape assert K == K2 c = torch.empty((M, N), device=a.device, dtype=TORCH_DTYPES[dtype]) grid = (num_programs, ) gemm_swizzle_kernel[grid]( a, b, c, M, N, K, a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1), BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k, SWIZZLE_MODE=swizzle_mode, GROUP_SIZE_M=group_size_m, DTYPE=TL_DTYPES[dtype], NUM_PROGRAMS=num_programs, ) return c# --- L2 hit-rate predictor (byte-granular LRU cache-residency simulation) ---class ByteLRUCache: """Byte-capacity LRU cache used to trace-simulate L2 tile residency. Each cached entry is identified by an arbitrary hashable key (here, an ('A'|'B', tile_index, k_index) tuple) and tracks its own byte size, so tiles of different shapes (A vs. B) can share one capacity budget. """ def __init__(self, capacity_bytes: int) -> None: self.capacity_bytes = capacity_bytes self.cache: "OrderedDict[Any, int]" = OrderedDict( ) # key -> size in bytes self.current_bytes = 0 self.hits = 0 self.misses = 0 def access(self, key: Any, size_bytes: int) -> None: if key in self.cache: self.hits += 1 self.cache.move_to_end(key) else: self.misses += 1 self.cache[key] = size_bytes self.current_bytes += size_bytes while self.current_bytes > self.capacity_bytes and len( self.cache) > 1: _, evicted_size = self.cache.popitem(last=False) self.current_bytes -= evicted_size @property def hit_rate(self) -> float: """Cumulative hit rate (percentage) across all access() calls so far.""" total = self.hits + self.misses return (self.hits / total * 100) if total > 0 else 0.0def predict_l2_hit_rate(fn: Callable[..., Tuple[int, int]], kwargs: Dict[str, Any], M: int, N: int, K: int, block_m: int, block_n: int, block_k: int, l2_cache_size_mb: float, concurrency: int, dtype_bytes: int) -> float: """Estimates the L2 cache hit rate for one swizzle mode via an offline, byte-granular LRU trace simulation, generalized to non-square A/B tile shapes. `fn` is a CPU pid-mapping function (e.g. linear_row_major, panel_swizzle) taking (pid, grid_m, grid_n, **kwargs) and returning (bx, by). Tile ids are grouped into waves of `concurrency` tiles (`min(total_tiles, measured_capacity)`) to approximate/reproduce real concurrent execution; within each wave, every tile's A/B sub-blocks are accessed for every K-iteration and replayed through an LRU cache sized to `l2_cache_size_mb`. This is a coarse model: it does not account for DRAM row-buffer locality, bank-level contention, or memory coalescing, so its output should be read as a directional signal, not a precise throughput predictor. """ grid_m = math.ceil(M / block_m) grid_n = math.ceil(N / block_n) grid_k = math.ceil(K / block_k) total_tiles = grid_m * grid_n cache = ByteLRUCache(l2_cache_size_mb * 1024 * 1024) a_tile_bytes = block_m * block_k * dtype_bytes b_tile_bytes = block_k * block_n * dtype_bytes tiles = [fn(pid, grid_m, grid_n, **kwargs) for pid in range(total_tiles)] for wave_start in range(0, total_tiles, concurrency): wave_tiles = tiles[wave_start:wave_start + concurrency] for k in range(grid_k): for bx, by in wave_tiles: cache.access(("A", bx, k), a_tile_bytes) cache.access(("B", k, by), b_tile_bytes) return cache.hit_ratedef rank_correlation(values_a: List[float], values_b: List[float]) -> float: """Spearman rank correlation with no external dependencies.""" def ranks(values: List[float]) -> List[float]: order = sorted(range(len(values)), key=lambda i: values[i]) r = [0.0] * len(values) for rank, idx in enumerate(order): r[idx] = rank return r ra, rb = ranks(values_a), ranks(values_b) n = len(values_a) mean_a, mean_b = sum(ra) / n, sum(rb) / n cov = sum((ra[i] - mean_a) * (rb[i] - mean_b) for i in range(n)) var_a = sum((x - mean_a)**2 for x in ra) var_b = sum((x - mean_b)** 2 for x in rb) if var_a == 0 or var_b == 0: return 0.0 return cov / math.sqrt(var_a * var_b)def _is_power_of_two(n: int) -> bool: """True if n is a positive power of two.""" return n > 0 and (n & (n - 1)) == 0AlgorithmEntry = Tuple[str, Callable[..., Tuple[int, int]], int, Dict[str, Any], int]def build_algorithm_list(group_sizes: List[int], grid_m: int, grid_n: int) -> List[AlgorithmEntry]: """Builds the list of (name, cpu_fn, swizzle_mode, kwargs, group_size_m) tuples to benchmark for the given grid shape. Always includes Linear Row-Major/Col-Major and one Grouped 2D Panel entry per value in `group_sizes`. Morton/Z-order and Bitwise XOR are included only when (grid_m, grid_n) is square and a power of two -- their bit tricks are not bijective otherwise (see module docstring), so they are skipped with a printed warning rather than silently producing incorrect results. """ algorithms = [ ("Linear (Row-Major)", linear_row_major, MODE_ROW_MAJOR, {}, 1), ("Linear (Col-Major)", linear_col_major, MODE_COL_MAJOR, {}, 1), ] for g in group_sizes: algorithms.append(( f"Grouped 2D Panel (g={g})", panel_swizzle, MODE_PANEL, { "group_size_m": g }, g, )) # Morton/XOR's bit tricks (interleaving, XOR-then-mod) are only bijective # (visit every tile exactly once) when the grid is square and a power of # two. Otherwise they silently skip some tiles and recompute others. if grid_m == grid_n and _is_power_of_two(grid_m): algorithms.append( ("Morton / Z-Order", morton_swizzle, MODE_MORTON, {}, 1)) algorithms.append(("Bitwise XOR", xor_swizzle, MODE_XOR, {}, 1)) else: print( f"[WARNING] Skipping Morton/XOR: grid {grid_m}x{grid_n} is not a " "square power-of-two, so those swizzles would drop/duplicate " "tiles and produce incorrect results.\n", file=sys.stderr, ) return algorithmsdef check_correctness(block_m: int, block_n: int, block_k: int, dtype: str = "fp16") -> None: """Verifies every swizzle mode against torch.matmul, on a small, fixed square power-of-two grid (so Morton/XOR are well-defined) that doesn't compete with the (possibly huge) benchmark tensors for VRAM. One program is launched per tile here (num_programs == total_tiles), so each program's loop runs exactly once regardless of --persistent. """ torch.manual_seed(0) grid_size = 4 # power of two, square: keeps every swizzle mode valid m, n, k = block_m * grid_size, block_n * grid_size, max(block_k * 2, 64) torch_dtype = TORCH_DTYPES[dtype] a = torch.randn((m, k), device="cuda", dtype=torch_dtype) b = torch.randn((k, n), device="cuda", dtype=torch_dtype) ref = torch.matmul(a, b) # fp32 tl.dot uses TF32 tensor cores by default (~10-bit mantissa, same # precision class as fp16), so it needs a loose tolerance too, with extra # margin for its worst-case (tail) accumulation error. rtol, atol = (1e-2, 1e-2) if dtype == "fp16" else (3e-2, 3e-2) num_programs = grid_size * grid_size algorithms = build_algorithm_list([1, 2, 4], grid_size, grid_size) for name, _, mode, _kwargs, group_size_m in algorithms: out = gemm_swizzle(a, b, mode, block_m, block_n, block_k, num_programs, group_size_m, dtype) torch.testing.assert_close(out, ref, rtol=rtol, atol=atol, msg=f"{name} produced incorrect results") del out del a, b, ref torch.cuda.empty_cache() print(f"Correctness check passed for all swizzle modes ({dtype}).\n")def run_benchmark(args: argparse.Namespace, hw_config: Dict[str, Any]) -> None: """Runs the full benchmark: correctness check, concurrency measurement, and one predicted-hit-rate + timed-throughput measurement per swizzle mode, printed as a summary table followed by a rank-correlation sanity check. """ check_correctness(args.block_m, args.block_n, args.block_k, args.dtype) torch.manual_seed(0) torch_dtype = TORCH_DTYPES[args.dtype] a = torch.randn((args.M, args.K), device="cuda", dtype=torch_dtype) b = torch.randn((args.K, args.N), device="cuda", dtype=torch_dtype) concurrent_capacity = measure_concurrent_capacity(a.device, args.block_m, args.block_n, args.block_k, args.dtype) grid_m = math.ceil(args.M / args.block_m) grid_n = math.ceil(args.N / args.block_n) total_tiles = grid_m * grid_n # Real achievable concurrency either way; the predictor uses this # regardless of launch style (see module docstring). concurrency = min(total_tiles, concurrent_capacity) num_programs = concurrency if args.persistent else total_tiles algorithms = build_algorithm_list(args.group_sizes, grid_m, grid_n) print("=" * 92) print("TRITON GEMM SWIZZLE BENCHMARK") print("=" * 92) print(f"Hardware Platform : {hw_config['gpu_name']}") print(f"Streaming Multiprocs: {hw_config['num_sms']} SMs") print(f"L2 Cache Capacity : {hw_config['l2_cache_size_mb']:.2f} MB") print(f"Matrix Dimensions : {args.M} x {args.N} x {args.K}") print(f"Tile Size : {args.block_m}x{args.block_n}x{args.block_k}") print(f"Data Type : {args.dtype}") print(f"Total Output Tiles : {total_tiles}") print(f"Measured Concurrent Capacity: {concurrent_capacity} CTAs") print("Kernel Mode : " f"{'Persistent' if args.persistent else 'Non-Persistent'}" f" ({num_programs} programs launched)") print("=" * 92) print(f"{'Algorithm':<28} | {'Pred. L2 Hit %':>14} | {'Time (ms)':>10} | " f"{'TFLOPS':>8}") print("-" * 92) hit_rates, tflops_list = [], [] for name, fn, mode, kwargs, group_size_m in algorithms: hit_rate = predict_l2_hit_rate(fn, kwargs, args.M, args.N, args.K, args.block_m, args.block_n, args.block_k, hw_config["l2_cache_size_mb"], concurrency, DTYPE_BYTES[args.dtype]) def bench_fn(mode=mode, group_size_m=group_size_m): return gemm_swizzle(a, b, mode, args.block_m, args.block_n, args.block_k, num_programs, group_size_m, args.dtype) ms = triton.testing.do_bench(bench_fn, warmup=args.warmup, rep=args.rep) tflops = (2 * args.M * args.N * args.K) / (ms * 1e-3) / 1e12 print( f"{name:<28} | {hit_rate:>13.2f}% | {ms:>10.3f} | {tflops:>8.2f}") hit_rates.append(hit_rate) tflops_list.append(tflops) print("-" * 92) corr = rank_correlation(hit_rates, tflops_list) print("Spearman rank correlation (predicted L2 hit rate vs measured " f"TFLOPS): {corr:.3f}") print( "(+1 = higher predicted hit rate always faster, -1 = always slower, 0 = uncorrelated)" )def _validate_args(args: argparse.Namespace) -> None: """Validates CLI arguments up front, raising a clear error instead of a cryptic Triton compile failure or silently-wrong behavior.""" for name in ("M", "N", "K", "warmup", "rep"): if getattr(args, name) <= 0: raise ValueError( f"--{name} must be positive, got {getattr(args, name)}") for name in ("block_m", "block_n", "block_k"): value = getattr(args, name) flag = name.replace("_", "-") if not _is_power_of_two(value): # tl.arange (used to build per-tile offsets) requires a power-of-two size. raise ValueError(f"--{flag} must be a power of two, got {value}") for g in args.group_sizes: if g <= 0: raise ValueError(f"--group-sizes values must be positive, got {g}")def main() -> None: """CLI entry point: parses arguments, resolves hardware config (auto-detected via detect_gpu_hardware, with optional CLI overrides), and runs the benchmark. """ hw_detected = detect_gpu_hardware() parser = argparse.ArgumentParser( description=( "Benchmark a Triton FP16/FP32 GEMM kernel across several " "thread-block swizzle orderings."), formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) parser.add_argument("--M", type=int, default=8192) parser.add_argument("--N", type=int, default=8192) parser.add_argument("--K", type=int, default=8192) parser.add_argument("--block-m", type=int, default=128) parser.add_argument("--block-n", type=int, default=128) parser.add_argument("--block-k", type=int, default=32) parser.add_argument( "--dtype", choices=["fp16", "fp32"], default="fp16", help="Data type shared by A, B, and C (fp32 uses the same tl.dot kernel path)", ) parser.add_argument( "--group-sizes", type=int, nargs="+", default=[1, 2, 4, 8, 16, 32, 64], help="Group sizes to try for the Grouped 2D Panel swizzle", ) parser.add_argument( "--persistent", action="store_true", help=("Launch exactly min(total_tiles, measured_capacity) " "persistent programs, each looping over several tiles, " "instead of one program per tile (see module docstring)"), ) parser.add_argument( "--num-sms", type=int, default=None, help="Override SM count reported/used by the predictor", ) parser.add_argument( "--l2-cache-mb", type=float, default=None, help="Override L2 cache size in MB used by the predictor", ) parser.add_argument("--warmup", type=int, default=25) parser.add_argument("--rep", type=int, default=100) args = parser.parse_args() _validate_args(args) if not torch.cuda.is_available(): raise RuntimeError("CUDA device required to run this benchmark.") final_sms = args.num_sms if args.num_sms is not None else hw_detected[ "num_sms"] final_l2 = (args.l2_cache_mb if args.l2_cache_mb is not None else hw_detected["l2_cache_size_mb"]) hw_config = { "num_sms": final_sms, "l2_cache_size_mb": final_l2, "gpu_name": hw_detected["gpu_name"], } run_benchmark(args, hw_config)if __name__ == "__main__": main()
|
It turns out that it is very difficult to correlate the actual kernel performance with the predicted L2 cache hit rate for both persistent and non-persistent GEMM kernel implementations.
For example, for $8192 \times 8192 \times 8192$ GEMM, with tile sizes of $128 \times 128 \times 32$, in spite of the lower L2 cache hit rate predicted for the row-major linear swizzle, there is no significant performance degradation compared to the grouped 2D panel swizzle.
What’s also interesting is that the column-major linear swizzle performs significantly worse than the row-major linear swizzle, despite having the same predicted L2 cache hit rate. In row-major order, row blocks of matrix B, which are contiguous in memory, are accessed in each wave. In column-major order, however, column blocks of matrix A, which are strided in memory, are accessed. Note that because both matrices are at least 16-byte aligned, the memory accesses with both row-major order and column-major order are still fully coalesced and vectorized. The reason to this performance difference is probably at much lower level, which we could not model, such as the DRAM access patterns.
|
1234567891011121314151617181920212223242526272829303132
|
$ python triton_gemm_swizzle_bench.py --M 8192 --N 8192 --K 8192 --block-m 128 --block-n 128 --block-k 32Correctness check passed for all swizzle modes (fp16).============================================================================================TRITON GEMM SWIZZLE BENCHMARK============================================================================================Hardware Platform : NVIDIA GeForce RTX 5080Streaming Multiprocs: 84 SMsL2 Cache Capacity : 64.00 MBMatrix Dimensions : 8192 x 8192 x 8192Tile Size : 128x128x32Data Type : fp16Total Output Tiles : 4096Measured Concurrent Capacity: 168 CTAsKernel Mode : Non-Persistent (4096 programs launched)============================================================================================Algorithm | Pred. L2 Hit % | Time (ms) | TFLOPS--------------------------------------------------------------------------------------------Linear (Row-Major) | 79.43% | 9.342 | 117.70Linear (Col-Major) | 79.43% | 11.548 | 95.21Grouped 2D Panel (g=1) | 79.43% | 9.383 | 117.18Grouped 2D Panel (g=2) | 79.52% | 9.486 | 115.90Grouped 2D Panel (g=4) | 85.57% | 9.476 | 116.03Grouped 2D Panel (g=8) | 92.13% | 9.464 | 116.18Grouped 2D Panel (g=16) | 95.27% | 9.482 | 115.95Grouped 2D Panel (g=32) | 88.06% | 9.485 | 115.92Grouped 2D Panel (g=64) | 79.43% | 11.636 | 94.49Morton / Z-Order | 89.39% | 9.590 | 114.66Bitwise XOR | 79.43% | 11.530 | 95.36--------------------------------------------------------------------------------------------Spearman rank correlation (predicted L2 hit rate vs measured TFLOPS): 0.045(+1 = higher predicted hit rate always faster, -1 = always slower, 0 = uncorrelated)
|
|
1234567891011121314151617181920212223242526272829303132
|
$ python triton_gemm_swizzle_bench.py --M 8192 --N 8192 --K 8192 --block-m 128 --block-n 128 --block-k 32 --persistentCorrectness check passed for all swizzle modes (fp16).============================================================================================TRITON GEMM SWIZZLE BENCHMARK============================================================================================Hardware Platform : NVIDIA GeForce RTX 5080Streaming Multiprocs: 84 SMsL2 Cache Capacity : 64.00 MBMatrix Dimensions : 8192 x 8192 x 8192Tile Size : 128x128x32Data Type : fp16Total Output Tiles : 4096Measured Concurrent Capacity: 168 CTAsKernel Mode : Persistent (168 programs launched)============================================================================================Algorithm | Pred. L2 Hit % | Time (ms) | TFLOPS--------------------------------------------------------------------------------------------Linear (Row-Major) | 79.43% | 9.574 | 114.85Linear (Col-Major) | 79.43% | 11.672 | 94.20Grouped 2D Panel (g=1) | 79.43% | 9.516 | 115.54Grouped 2D Panel (g=2) | 79.52% | 9.638 | 114.08Grouped 2D Panel (g=4) | 85.57% | 9.663 | 113.79Grouped 2D Panel (g=8) | 92.13% | 9.636 | 114.11Grouped 2D Panel (g=16) | 95.27% | 9.631 | 114.17Grouped 2D Panel (g=32) | 88.06% | 9.869 | 111.41Grouped 2D Panel (g=64) | 79.43% | 11.875 | 92.59Morton / Z-Order | 89.39% | 9.703 | 113.31Bitwise XOR | 79.43% | 11.681 | 94.13--------------------------------------------------------------------------------------------Spearman rank correlation (predicted L2 hit rate vs measured TFLOPS): 0.064(+1 = higher predicted hit rate always faster, -1 = always slower, 0 = uncorrelated)
|
Such performance differences might diminish if other benchmark configurations are used or a different compute platform is considered.
#
Predicting how thread block swizzle affects the actual performance of GEMM kernels is challenging. It might just be more pragmatic to empirically benchmark different swizzle strategies for specific kernels and configurations rather than relying solely on predicted L2 cache hit rates.
The order of thread block execution can significantly impact cache utilization and overall performance. In some other non-GEMM applications, sometimes simply switching between the row-major and column-major execution order can lead to noticeable performance differences.
#
CUDA Thread Block Swizzle
source & further reading
leimao.github.io — original article
CUDA Device Max Connections
Residual-Quantized Variational Autoencoder
CUDA Multi-Process Service