Data layout and memory access
Write a complete coordinate-to-address contract and decide whether conversion or persistent packing improves the full producer-consumer path
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.
- Ability to print tensor shape, stride, datatype, and storage offset
- Previous section's coalescing model
- Introductory understanding that packed values require encoding metadata
The recurring 7-billion-parameter inference model runs on one 80 GiB GPU. A profiler shows that one quantized projection kernel is fast after its weights are packed, but request startup now contains a 40-millisecond packing operation and a later debugging operator cannot read the packed tensor. A memory report shows both ordinary and packed weight copies resident. The local matrix speedup has created a representation decision for the complete model server.
Inspect the tensor with framework shape, stride, datatype, and storage-offset APIs, then capture the packing and consumer kernels with PyTorch Profiler or Nsight Systems. Shape describes logical positions. Strides and the packing specification explain how those positions map to byte addresses. The trace records when conversion actually moves bytes; a source-level reshape or transpose name does not establish whether the operation was only a view.
A tensor's shape gives the number of logical positions along each dimension. Its layout defines how those positions map to memory addresses. For a simple contiguous row-major matrix, moving one column changes the address by one element and moving one row changes it by the row width. A transposed view can preserve the same storage and change only the strides used for indexing.
GPU kernels often use layouts more complex than row-major and column-major. A matrix library can interleave or swizzle tiles to match shared-memory banks and matrix instructions. Quantized weights can pack several low-bit values into one word and store a scale for each group. A key-value (KV) cache can place logical token blocks in non-contiguous physical pages. Each format needs an explicit coordinate-to-address rule.
Layout affects more than one kernel. A packed weight format can make a general matrix multiplication (GEMM) fast but require conversion before another operator. Converting back after every use can erase the gain. Preserving the packed form can save repeated conversion but add compatibility, memory, and fallback requirements. The correct scope is the complete producer-consumer region.
Define the layout contract
A complete layout description includes logical shape, physical strides, blocking, padding, element encoding, and auxiliary metadata. For packed low-bit weights, the contract must define which bits represent each logical value and where scales or zero points reside. For blocked matrices, it must define the order of blocks and the order of elements within a block.
Separate the logical tensor from its physical representation. Mathematical operators should be specified in logical coordinates. Target-specific kernels then implement those semantics for one or more layouts. This separation permits a correct general implementation, architecture-specific fast paths, and tests that compare their outputs.
Views and conversions are different. A view changes shape or strides without moving bytes when the existing storage can represent the new indexing. A contiguous conversion, transpose kernel, or packing operation moves and possibly transforms bytes. Framework syntax can make both operations look inexpensive, so verify whether storage is materialized.
Calculate conversion amortization
Suppose packing weights costs 40 ms and the packed GEMM saves 18 microseconds per invocation. Ignoring other effects, the break-even count is 40 ms divided by 18 microseconds, or about 2,223 invocations. A long-running server may exceed that count quickly. A short command-line job may never recover the setup cost.
The calculation changes when packed weights are stored in the model artifact. Startup packing disappears, but artifact size, version compatibility, loading code, and fallback behavior enter the system contract. Keeping both packed and ordinary copies removes some conversion cost but consumes more memory. Evaluate cold start, steady state, and memory capacity separately.
Activation layouts usually have a shorter lifetime than weight layouts. A per-request conversion must be amortized across fused or layout-aware consumers in the same request. Count the bytes read and written by each conversion, the number of later operations that benefit, and any padding added to the working set.
Layout can change allocation and scheduling
Paged KV storage maps logical sequence positions to physical blocks. This adds address indirection inside attention, but it lets the allocator place variable-length sequences without requiring one contiguous region for each request. The serving scheduler can reuse freed blocks and batch requests with different histories more flexibly.
Distributed layouts make another system-level choice. Sharding one tensor dimension can make local matrix operations efficient while requiring communication before the next operator. The physical arrangement therefore encodes both local memory access and inter-device traffic. Select it by analyzing the operator graph, not one kernel in isolation.
Target-specific layouts need explicit boundaries. Record supported shapes, alignment, architecture, datatype, and fallback. Benchmark common aligned shapes and adversarial tails. A format that is fast only for one tile family can still be valuable, but its dispatch conditions must prevent silent misuse.
Worked example: Deciding whether to preserve packed weights
A quantized GEMM uses a hardware-friendly packed weight format. Packing takes 40 ms, while each invocation saves 18 microseconds compared with the ordinary layout.
- Compute the simple break-even: 40 ms / 18 µs is about 2,223 invocations. A long-lived serving replica will easily amortize this; a one-shot job may not.
- Add storage and loading behavior. If packed weights are serialized in the model artifact, startup packing disappears but artifact compatibility and versioning become part of the contract.
- Trace consumers. Debugging, fine-tuning, or fallback operators may need ordinary weights. Decide whether to keep both copies, unpack on demand, or provide packed-aware paths.
- Benchmark cold start, steady-state inference, and memory footprint. The decision depends on deployment lifetime and capacity, not only GEMM speed.
Conclusion: The 2,223-invocation break-even is the first decision point, not the final result. Deployment lifetime, memory capacity, artifact format, mutation frequency, and fallback requirements determine whether preserving the packed representation improves the complete system.
Common errors
- Treating a transpose as free framework bookkeeping. If strides cannot represent the desired view, bytes must move somewhere.
- Benchmarking only the consumer that benefits. Conversion, padding, and fallback paths belong in the same accounting boundary.
- Assuming an exotic layout is portable. Keep semantics and a correct general path separate from target-specific physical representation.
Reference summary
A layout maps logical coordinates to physical addresses. The contract includes shape, strides, blocking, padding, element encoding, and any required metadata. Matrix swizzles, packed low-bit weights, and paged KV storage require more than a row-major or column-major label.
Evaluate layout across producers and consumers. A conversion is beneficial only when later savings repay its time, traffic, memory footprint, and compatibility cost. Long-lived weights can amortize one-time packing more easily than per-request activations.
- Trace every consumer before adding a transpose or packing kernel.
- Include padding and metadata overhead in the memory model.
- Specialize common shapes while keeping a correct general path.
- Measure conversion amortization: setup cost divided by number of reuses.
- Treat KV-cache layout as a serving-scheduler decision, not only an attention-kernel detail.
Knowledge check
When is it rational to keep weights in a packed hardware-specific layout even though ordinary PyTorch operations cannot read it directly?
Show the answer guide
- Persistent packed weights can be rational when one conversion at model load is reused by enough later operations to repay the packing time and extra stored bytes.
- The packed layout must match the selected matrix instructions and every production consumer, or the path needs explicit conversion for incompatible operations.
- The decision requires a break-even invocation count, cold-start and steady-state timing, memory footprint, shape coverage, and a fallback for unsupported consumers.
Calculate layout amortization
Compare on-demand packing with one-time persistent packing for a declared weight matrix and consumer sequence. Measure packing, packed-operation, ordinary-operation, and fallback conversion times.
Evidence to keep: Producer-consumer layout diagram, raw timings, memory footprint, compatibility table, break-even invocation count, and the assumed replica lifetime.
The memory model predicts execution time and resource limits. Profiling tests those predictions with application, kernel, and instruction measurements.
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.