Transformer-Based Resource and Stage-Aware Scheduling for Model-Parallel LLM Inference

Research Note

When a large language model does not fit on a single GPU, it has to be split across nodes, and a scheduler decides which work goes where. Existing schedulers tend to carry three assumptions. They treat inference as a uniform workload, they assume the nodes are identical, and they are designed for data centres where interconnect bandwidth is plentiful, so network congestion never enters the picture.

Inference is in fact made of two phases with different profiles. Prefill processes the whole input at once and is compute-heavy, whereas decode emits one token at a time and is light on computation but sensitive to communication round trips. Scheduling both as though they were the same kind of work forces a compromise that suits neither.

The proposed scheduler models heterogeneous GPU clusters, communication costs and per-stage characteristics explicitly. Three approaches are compared, a heuristic stage-based policy, a continuous-batching baseline in the style of vLLM, and a transformer-based scheduler trained by imitation to reproduce and then improve on the heuristic.

Across eight scenarios covering oversized models, prefill-dominant and mixed workloads, heterogeneous and bandwidth-limited clusters, strict SLO constraints and multi-tenant deployments, the learned scheduler cut latency by up to 50% under bandwidth-constrained or heterogeneous conditions while keeping throughput within 20 to 30% of vLLM, and improved latency by 3 to 17% over its heuristic teacher while preserving full feasibility. Continuous batching remained better on high-bandwidth fabrics, which identifies bandwidth as a first-order determinant of the right scheduling strategy.