CUDA Thread Block Swizzle

A technical blog post analyzes how CUDA thread block swizzle algorithms change L2 cache residency and GEMM kernel performance by remapping program IDs to tile locations. The post compares four swizzle schemes on an 8×8 GEMM tiled grid — row-major linear, grouped 2D panel swizzle with group size 4, Morton/Z-order swizzle, and bitwise XOR swizzle — noting that Morton/Z-order and bitwise XOR swizzles only produce bijective mappings for grid sizes that are powers of two. The author simulates L2 cache hit rates with an offline model and benchmarks GEMM kernels under each swizzle, identifying grouped 2D panel swizzle as the most common choice in practice apart from the usually default linear swizzle.

CUDA Thread Block Swizzle 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 https://leimao.github.io/blog/CuTe-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 https://dev.to/sachin tolay 052a7e539e57/understanding-dram-internals-how-channels-banks-and-dram-access-patterns-impact-performance-57ng . | 1234567891011121314151617181920212223242526272829303132 | bash $ 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 | bash $ 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