Library selection and shape families
Build and validate a GEMM dispatch policy for a weighted shape distribution
Before you begin
Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.
- Benchmarking
- Matrix-instruction constraints
An inference server records matrix calls from one hour of traffic. Prompt processing produces matrices with thousands of token rows, while token generation produces matrices with one to 128 rows depending on the active request batch. A cuBLASLt or CUTLASS benchmark table shows that the fastest kernel changes across those shapes. One implementation can reach high tensor throughput for M=4096 yet waste most output-tile lanes for M=1.
The runtime decision that maps a declared matrix shape, datatype, layout, and output operation to an implementation is a dispatch policy. A useful investigation captures the real shape histogram, benchmarks candidate implementations with the actual bias or activation epilogue, and weights each result by how often the shape occurs. The output is a table of measured regions and a fallback, not one supposedly universal kernel.
There is no universally fastest GEMM kernel. A schedule tuned for a huge square training matrix may be wasteful for a batch-one decode projection; a persistent kernel may shine when repeated weights fit a reuse strategy; grouped GEMM may serve irregular experts better than thousands of tiny launches. Production performance depends on dispatching the right implementation for a distribution of shapes.
Mature libraries combine handwritten or generated kernels, architecture-specific features, heuristics, and autotuning. The engineering task is not to outguess them casually. The task is to define the workload family, measure dispatch behavior, identify gaps, and decide whether configuration, packing, shape changes, or a specialized kernel earns its maintenance cost.
A workload family is more than a list of dimensions. The family includes how often each shape occurs, whether calls are latency-critical, which epilogue is required, how weights are stored, whether workspace is available, and whether first-use compilation or tuning is permitted. The operating conditions can change the best implementation even when M, N, and K are identical.
Shape regimes change the available parallel work
Large M and N create many independent output tiles, so the GPU can run many thread blocks. A batch-one decode projection may have M=1 or a small active-token count. The decode shape exposes far fewer output tiles and may read a large weight matrix for little reuse. A grouped mixture-of-experts workload contains many matrices of different M values, so assigning one conventional launch to each group can create launch overhead and load imbalance.
K changes how long each output tile remains in the main loop. Large K amortizes setup and epilogue cost but keeps accumulators live for longer. Small K can make the epilogue and stores dominant. Ragged dimensions create partial tiles. Record transpose state, leading dimensions, alignment, datatype, and epilogue because these values restrict the valid kernel set.
Heuristics, autotuning, and dispatch
A heuristic predicts a short list of likely kernels from the problem and hardware. Selection overhead is low, but the model can miss unusual shapes. Autotuning executes candidates and records the fastest measured result. Autotuning can find cases that the heuristic misses, but compilation, benchmarking, and cache storage are real costs.
An online service should not launch an unbounded tuning search inside a latency-critical request. The service team can tune offline, tune during controlled warm-up, or use bounded background tuning with a safe initial choice. Cache keys must include every property that changes correctness or schedule quality. A cache keyed only by M, N, and K is incorrect if datatype, layout, epilogue, or target architecture can differ.
Changing the surrounding representation
Kernel selection is not restricted to choosing among existing calls. Packing a constant weight matrix can improve vectorized access. Padding K to an instruction multiple can enable a much faster matrix path, provided the extra arithmetic and storage are smaller than the gain. Batching several small matrices can reduce launch overhead. A grouped kernel can schedule irregular expert matrices in one launch.
Every representation change has a lifecycle. A packed weight may be reused across thousands of requests, which makes one-time packing inexpensive. Packing a transient activation for one call may cost more than it saves. A standalone GEMM benchmark that begins after packing excludes this cost. End-to-end comparison must start before the conversion and end after the required output layout is available.
Evaluating a dispatch policy
Collect a histogram of real calls rather than selecting a few memorable shapes. Weight each measurement by call frequency for throughput work, by time contribution for optimization prioritization, or by request class for latency objectives. Report tail shapes even if they contribute little average time because they can affect service latency or fallback reliability.
A dispatch table should use stable regions instead of one rule per observed dimension when possible. The table also needs a general fallback and versioned benchmark evidence. Re-run the distribution after a GPU, library, compiler, model architecture, or batching-policy change. The winning regions can move even when the dispatch source code is unchanged.
Worked example: Building a projection dispatch table
A decoder serves a mix of prefill and decode. Its projections share N and K but M ranges from one active token to tens of thousands of prompt tokens.
- Collect a weighted histogram of M under representative load. Separate decode microbatches, chunked prefill, and large offline prefill rather than averaging them.
- Record N, K, datatype, weight layout, alignment, epilogue, and concurrency with each M bucket. Do not merge calls that have different kernel eligibility.
- Benchmark candidate library algorithms with the actual epilogue, alignment, precision, and packed weight format. Include warm cache and realistic concurrency if they affect selection.
- Define shape regions with stable winners and a correct fallback. Avoid one rule per exact M unless cache and maintenance costs are justified.
- Replay the full distribution and compare weighted latency, throughput, workspace, and startup behavior. Revalidate when the GPU, library, or model shape changes.
- Keep the complete general library path as a fallback and add boundary tests on both sides of every dispatch threshold.
Conclusion: Dispatch applies kernel selection to the workload distribution. Evaluate the dispatch policy across the complete distribution instead of one shape.
Common errors
- Benchmarking one square shape and generalizing to all matrix work. Aspect ratio and granularity create different regimes.
- Autotuning inside a latency-critical request. Tuning cost and cache misses belong in the deployment design.
- Ignoring library maturity. A custom kernel should beat the appropriate baseline and justify correctness and maintenance, not only beat a naive loop.
- Reporting a packed-kernel time without the packing lifecycle. Reuse count determines whether the conversion is economically valid.
Reference summary
Production GEMM systems dispatch among implementations. Large square training matrices, grouped mixture-of-experts work, batched small matrices, and narrow decode projections favor different tiles and scheduling. Vendor libraries and CUTLASS profilers search a large space; your task is to characterize the actual shape family and include dispatch overhead, workspace, and epilogue requirements.
- Record M, N, K, batch/group count, datatype, layout, alignment, and epilogue.
- Weight measurements by production frequency or total time contribution.
- Separate first-run autotuning/compilation from steady state.
- Keep a fallback for uncommon shapes and future architectures.
- Compare against mature libraries before you claim a performance improvement.
Knowledge check
Why might one GEMM kernel be fastest for large square matrices but slower for the model’s decode workload?
Show the answer guide
- A large square GEMM exposes many output tiles, has a long reduction dimension, and supplies enough work to amortize launch, setup, and pipeline costs. A kernel with large tiles and several pipeline stages can use tensor cores efficiently in that regime.
- A decode GEMM often has M equal to one or a small batch of token rows while N and K remain large. The same large-tile kernel can leave most rows inactive, launch too few blocks, or spend too much time on staging and synchronization.
- Decode can instead favor narrow tiles, split-K reduction, persistent scheduling, fused neighboring operations, or a different library algorithm. Datatype, alignment, epilogue, grouped shapes, and parallel topology also change the winning path.
- Kernel selection must therefore use the model's actual shape distribution and service objective, not performance on one large square benchmark.
Build a shape-aware GEMM dispatch table
Use cuBLASLt, CUTLASS, or another tunable GEMM library to benchmark at least four algorithms on FP16 shapes (4096,4096,4096), (2048,4096,4096), (32,4096,4096), (8,4096,4096), and (1,4096,4096), where tuples are M,N,K. Include the model's required bias or activation epilogue if applicable and verify results against FP32 accumulation.
Evidence to keep: Submit the benchmark harness, correctness results, a matrix of latency and achieved throughput by shape and algorithm, and the resulting dispatch table. Include one profiler comparison explaining why the square-shape winner loses on M=1 or M=8.
Efficient GEMM reuses input tiles for dense arithmetic. Fusion and attention can also eliminate writes and reads of intermediate tensors.
Sources and further reading
The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.