Specialisation & Co-design

Specialisation and co-design pull in opposite directions: one separates responsibilities, and the other couples design decisions. However, in modern ML systems, they work together to realise gains from specialisation beyond any individual component. While specialisation generally exploits workload differences, co-design coordinates the resulting choices to improve performance against the system’s objective. Specialised components require careful boundary selection, interface design, and joint optimisation.

Different compute and memory requirements for [encoding,] prefill and decode stages in modern inference systems inform disaggregated serving with separate GPU pools. Similarly, attention and FFN layers have different locality and parallelism requirements; RL rollout generation and training have different execution patterns.

Work separation lets us optimise and scale individual components, but it also introduces interfaces the system must manage, adding transfers, queues, synchronisation, and stranded capacity. Boundaries between components can be physical, logical, temporal or representational, so a different execution strategy does not necessarily suggest a distributed system.

The resulting system’s performance should justify the separation and boundary-management overhead by improving the global system’s objective under product constraints. Optimising individual components of the disaggregated inference stack independently without accounting for KV transfer and joint latency budgets would likely introduce unnecessary bottlenecks. Modality mix, cache state, prompt and output length distribution, arrival patterns, hardware and network topology all influence efficient provisioning. Even as execution becomes more distributed, design decisions remain tightly coupled.

Model design adds another layer of complexity and further shapes deployment infrastructure. Conversely, existing system bottlenecks influence model architectures and change production economics. GQA and MLA change KV cache requirements and communication and allocation patterns in serving clusters. As such, desired memory pressure and communication traffic become an architectural choice made at the model design stage.

State management places locality and bandwidth constraints on the system within the scale-up domain and beyond. Routing strategy in MoE models often accounts for heterogeneous networking bandwidth. KV state becomes a distributed scheduling and storage resource. Even with local execution, FlashAttention uses tiling and, later, warp specialisation to overlap data movement and computation.

Interfaces between specialised components determine levels of independence and enable efficient communication. Prefill and decode workers coordinate access to transferred KV state. Rollout workers and trainers can coordinate to enable staleness-aware training. PagedAttention borrows operating systems’ virtual memory abstraction to connect attention kernels to the KV-cache allocator.

With the above in mind, a framework for co-designing a specialised system could be:

Define objective and constraints (e.g., throughput under latency targets)
Diagnose bottleneck (e.g., exposed communication or queueing)
Identify opportunity for specialisation (e.g., different resource demands)
Choose boundary and interface (e.g., state placement and transfer format)
Optimise jointly (e.g., adjust batches or pool capacity)
Evaluate against objective and constraints

Specialisation & Co-design

Relevant publications