Tensor layout optimization for prefill-decode disaggregation requires tailoring memory access patterns and kernel scheduling to each phase's distinct characteristics: prefill is compute-intensive and benefits from tensor parallelism with fused kernels, while decode is memory-bandwidth-bound and requires careful KV cache management and cache-aware scheduling. Consumer GPU deployment of disaggregated pipelines demands hardware-aware kernel fusion and intra-GPU resource allocation strategies that avoid memory bottlenecks while maintaining efficient data transfer between pipeline stages.
The separation of prefill and decode phases in large language model (LLM) inference represents a paradigm shift in optimization strategy, fundamentally driven by the divergent resource demands of each stage [1][5]. Prefill—the processing of input prompt tokens—demands sustained computational throughput, while decode—the generation of output tokens—bottlenecks on memory bandwidth and latency [8][17]. Tensor layout optimization and hardware-aware kernel scheduling emerge as critical techniques for realizing the performance benefits of disaggregated inference on consumer-grade GPUs, where memory constraints are more severe than enterprise systems [2][3].
The architectural foundation for optimization begins with understanding inherent phase differences. Prefill operates on large batch dimensions with substantial compute-to-memory ratios, making it amenable to parallelization across multiple GPUs through tensor parallelism [5]. Decode, conversely, processes a single token at a time per sequence, creating severe memory-bandwidth bottlenecks despite requiring minimal computation [1][8]. This asymmetry enables independent scaling: prefill systems can optimize for arithmetic intensity, while decode systems can optimize for memory throughput [6].
KV cache management exemplifies this distinction. Traditional architectures store KV caches in GPU memory to accelerate decoding operations [2], but the actual bottleneck is not computation but rather the bandwidth required to read massive cached tensors. For disaggregated pipelines, KV cache must be efficiently transferred between prefill and decode stages, introducing data movement costs that tensor layout decisions directly influence [4][7].
Tensor layout—the arrangement of multi-dimensional data in linear memory—profoundly impacts both computation efficiency and inter-stage transfer overhead. For prefill, standard dense row-major or column-major layouts remain appropriate because the compute kernel (matrix multiplication) exhibits high arithmetic intensity [5]. However, the KV cache produced by prefill must be reformatted for efficient decode consumption.
Optimal decode-stage tensor layouts prioritize sequential memory access patterns that align with GPU cache line sizes and memory coalescing requirements. Rather than storing KV tensors in their natural dimension order, disaggregated systems benefit from layouts that group head dimensions contiguously, enabling efficient loading during attention computation [2][11]. This is particularly critical on consumer GPUs where L2 cache is limited (typically 4-6 MB) and memory bandwidth does not scale linearly with compute capability.
The transfer of KV cache between nodes or GPUs introduces another layout consideration: the transfer format should minimize serialization overhead while maintaining layout compatibility with the decode kernel. Standard approaches use contiguous memory layouts for serialization, but hardware-aware optimizations may employ compression or tiling schemes that reduce transfer volume [7][19].
Kernel fusion—combining multiple operations into a single kernel—emerges as a foundational technique for both prefill and decode optimization [11][12]. In prefill, fusing normalization, linear projections, and other operations reduces intermediate tensor materialization and improves cache locality [12]. FlashAttention exemplifies this approach, fusing all attention computation steps into an IO-aware kernel that tiles computation and maintains working sets in fast on-chip SRAM rather than expensive HBM [13].
For decode phases, kernel fusion is equally critical. Fused kernels for attention, normalization, and projection operations dramatically reduce memory bandwidth requirements by avoiding redundant reads of KV caches and input activations [14][15]. On consumer GPUs with limited memory bandwidth (approximately 500-900 GB/s for current generation), kernel fusion can provide 2-4x throughput improvements compared to unfused implementations.
The scheduling of fused kernels across GPU hardware requires awareness of SM (streaming multiprocessor) occupancy and warp scheduling. With multiple independent decode sequences executing concurrently, kernel schedulers must balance work distribution to avoid register pressure or memory bank conflicts [6].
Effective cache utilization distinguishes high-performance disaggregated systems from naive implementations. L1 caches on consumer GPUs typically provide 128 KB per SM, while L2 caches provide 4-6 MB shared across all SMs. Prefill kernels working on full batches can exploit L2 caches effectively through tiling strategies: processing blocks of sequences enables temporal reuse of linear layers and attention masks [11].
Decode phases present different cache challenges. Each sequence's KV cache grows during generation, but at decode time, the entire historical KV tensor must reside in memory. Tiling strategies that partition the KV cache across decode iterations minimize L2 cache pressure while ensuring compute kernels can access required data within a single SM's local memory view [2][4].
For disaggregated pipelines specifically, cache efficiency must account for inter-stage data movement. When prefill and decode run on separate GPUs, PCIe or network transfer introduces additional latency. Consumer systems with PCIe 4.0 (16 GB/s) face severe bandwidth constraints compared to enterprise interconnects (100+ GB/s), necessitating aggressive tensor layout optimization to minimize transfer sizes [7][16].
Kernel scheduling in disaggregated systems must account for dynamic load imbalance between prefill and decode stages. Prefill latency varies with input prompt length, while decode duration depends on output sequence length and batching factors [9]. Task schedulers must overlap prefill execution on idle resources with ongoing decode sequences, avoiding idle GPU stalls [3][8].
Intra-GPU disaggregation—partitioning a single GPU's SMs between prefill and decode stages [6]—introduces additional scheduling complexity. With limited SM count on consumer GPUs (typically 80-128 SMs), dividing resources requires careful profiling to determine optimal allocation. Static allocation may leave prefill SMs idle during low-throughput periods, while dynamic allocation incurs context switching overhead.
Hardware-aware schedulers leverage GPU performance counters to monitor SM occupancy, memory bandwidth utilization, and cache hit rates, dynamically adjusting kernel launch dimensions and resource allocation [6]. This approach requires tight integration with kernel implementations that support variable thread block granularity and cooperative grid launch features.
In disaggregated architectures using separate GPUs or nodes, KV cache transfer represents a critical bottleneck [4][7][10]. Full disaggregation transfers complete KV tensors over network interfaces, while intra-GPU disaggregation transfers within GPU memory hierarchies [6][16].
Layout conversion costs dominate transfer overhead on consumer systems. If prefill produces KV tensors in one layout optimized for batch matrix multiplication, and decode requires a different layout optimized for memory coalescing, conversion kernels must execute before transfer. Optimizing this conversion—potentially through in-place transformations or clever layout design that satisfies both phases—yields substantial performance gains [2][19].
Abstraction layers like those provided by AWS and other cloud providers handle transfer method selection (libfabric, UCCL, GPUDirect Storage) [19], but tensor layout decisions ultimately determine transfer efficiency. Layouts that minimize per-tensor metadata and enable streaming transfers without intermediate buffering perform best on consumer hardware with limited auxiliary memory.
Consumer GPUs present distinct constraints compared to enterprise accelerators. Limited memory (8-24 GB vs 40-80 GB on data center GPUs) restricts batch sizes during prefill and forces aggressive KV cache quantization or eviction strategies. Disaggregation can partially mitigate this by separating compute demands: prefill can operate with smaller batches while decode focuses on throughput [3][8].
Memory fragmentation becomes problematic with disaggregated pipelines, as prefill allocates large temporary buffers for batch processing while decode maintains persistent KV caches. Kernel fusion reduces intermediate allocations, while careful memory pooling strategies prevent fragmentation [11][14].
Power constraints on consumer systems are less severe than thermal constraints. Kernel fusion and optimized cache utilization reduce memory bandwidth demand, lowering overall power consumption and enabling sustained high-frequency operation without thermal throttling.
Tensor layout optimization for prefill-decode disaggregation on consumer GPUs requires a multifaceted approach: careful layout design that serves both phases, aggressive kernel fusion to reduce memory pressure, cache-aware scheduling that exploits limited on-chip resources, and dynamic kernel scheduling that balances computational phases [10][17][18]. The fundamental asymmetry between compute-intensive prefill and memory-intensive decode justifies disaggregation, while hardware awareness ensures that the inherent optimization potential translates to measurable performance improvements despite consumer GPU constraints.