# Topology-Aware GPU scheduling with KAI Scheduler - URL: https://hrittikhere.com/posts/kubernetes-topology-aware-scheduling-kai - Published: 2026-07-12 - Category: Kubernetes - Tags: KAI Scheduler, Kubernetes, GPU scheduling, gang scheduling, topology-aware scheduling, distributed training - Read time: 15 MIN READ How KAI Scheduler uses topology-aware placement and gang scheduling to keep distributed GPU training inside one region, zone, or rack on Kubernetes. ## Key Takeaways - Kubernetes will happily split a distributed GPU job across regions unless something tells it the pods are one tightly coupled workload. - Topology-aware placement plus gang scheduling let you require the complete job to fit inside one region, zone, or rack, or stay pending. - Fast networking like NVLink and InfiniBand only helps once placement puts the communicating GPUs inside the same accelerated domain. - Inference benefits too: disaggregated prefill and decode serving stays fast when placement keeps the KV-cache transfer local and each tensor-parallel group on one rack or NVLink domain, giving steadier tail latency and lower cost per request. ## FAQs ### What is topology-aware scheduling in Kubernetes? Topology-aware scheduling places the pods of a workload based on the physical and network layout of the cluster, such as regions, zones, racks, and hosts, instead of treating every node as interchangeable. KAI Scheduler models that layout with a cluster-scoped Topology resource built from node labels and constrains placement with required and preferred levels. ### Why do distributed GPU jobs need gang scheduling? Distributed training pods synchronize constantly, so a partially started job either blocks on missing peers or runs with a degraded communication path. Gang scheduling admits the whole group together or keeps it pending together, which also makes topology constraints apply to the complete job rather than to each pod individually. ### Does InfiniBand remove the need for topology-aware placement? No. InfiniBand and RoCE accelerate communication inside a data center or cluster fabric, but they do not connect distant regions with one lossless, microsecond-latency network. Placement decides whether a workload can use the fast fabric at all, and the fabric then decides how fast communication inside that placement is. ## Article You have four NVIDIA A100 40 GB GPUs in one Kubernetes cluster. Two are attached to nodes in `us-central1` in Iowa and two are attached to nodes in `asia-northeast3` in Seoul. You submit a distributed training job that needs two GPUs. Kubernetes starts both pods, so everything looks healthy at first. Twenty minutes later, GPU utilization is low and the loss curve is moving slowly. Then this command reveals the problem: ```bash kubectl get pods -o wide ``` One training pod is running in Iowa. The other is running in Seoul. Every collective operation between them is crossing the Pacific. Kubernetes did not split the job across regions on purpose. The problem is that nothing told the scheduler that these two pods form one tightly coupled workload and must remain close to each other. KAI Scheduler solves this with [topology-aware scheduling](https://github.com/NVIDIA/KAI-Scheduler/blob/main/docs/topology/README.md) and [gang scheduling](https://github.com/NVIDIA/KAI-Scheduler/blob/main/docs/batch/README.md). Together, they let you express a simple rule: > Schedule the complete distributed job inside one region, or keep the complete job pending. This article explains why that rule matters, how network distance affects GPU workloads, where InfiniBand and NVLink help, and how to configure KAI Scheduler without accidentally creating constraints that enforce nothing. ## Why distance matters for distributed GPU workloads A distributed training job does not use multiple GPUs as independent workers. The GPUs exchange data throughout the training step. With data parallel training, each GPU computes gradients from a different batch of data. The gradients must then be combined, commonly through an `all-reduce` collective, before the next step can continue. Tensor parallelism can communicate even more frequently because GPUs exchange intermediate tensors during model execution. The slowest communication path can therefore delay every participant. ### A practical comparison The following values are representative orders of magnitude, not guaranteed measurements. Actual performance depends on the GPU model, machine type, NIC, topology, congestion, collective library, and cloud network. | Placement | Typical RTT scale | Link bandwidth scale | What it means | | --- | ---: | ---: | --- | | GPUs inside one NVLink/NVSwitch system | Microseconds | Hundreds of GB/s | Best for tensor parallelism and frequent collectives | | Nodes on the same high-speed cluster fabric | Tens of microseconds | 100 to 400 Gb/s per NIC | Suitable for multi-node training when RDMA is configured correctly | | Nodes in one cloud zone over Ethernet | Sub-millisecond to a few milliseconds | Often tens to hundreds of Gb/s | Can work, but depends heavily on the instance and network tier | | Nodes in different regions | Tens to hundreds of milliseconds | WAN-limited and variable | Usually unsuitable for synchronous GPU collectives | | Iowa to Seoul | Roughly 130 to 200 ms RTT is a reasonable planning range | Far below local GPU fabric performance | Each synchronization can become an intercontinental round trip | For the real value in your environment, measure the VM-to-VM path. [Google Cloud's Performance Dashboard](https://docs.cloud.google.com/network-intelligence-center/docs/performance-dashboard/how-to/view-google-cloud-latency) exposes regional latency data, and tools such as `ping`, `iperf3`, NCCL Tests, and NVIDIA's network diagnostics help measure the path seen by the workload. ### A simple latency example Assume a training step performs 100 synchronization points and the cross-region round-trip time is 150 ms. If each synchronization waits on even one network round trip, latency alone can contribute: ```text 100 synchronizations × 150 ms = 15 seconds ``` That is before transferring the gradients. Now assume each GPU must exchange a 4 GB gradient payload. On a theoretical 100 Gb/s link, which equals 12.5 GB/s before protocol overhead, moving 4 GB takes at least: ```text 4 GB ÷ 12.5 GB/s = 0.32 seconds ``` Real collective throughput will be lower because of protocol overhead, contention, routing, and the collective algorithm. Across regions, available throughput may also vary. The total communication time becomes a combination of: ```text communication time ≈ transfer time + latency cost + synchronization overhead ``` This is why adding more GPUs can sometimes make a distributed job slower. Compute capacity increased, but communication became the bottleneck. ## Where NVLink, NVSwitch, and InfiniBand help Topology-aware scheduling does not replace fast networking. It makes sure the workload is placed where fast networking is available. ### NVLink and NVSwitch NVLink provides high-bandwidth GPU-to-GPU communication inside supported systems. NVSwitch connects multiple GPUs into a larger high-bandwidth fabric. On the [A100 platform](https://www.nvidia.com/en-us/data-center/a100/), third-generation NVLink delivers 600 GB/s of GPU-to-GPU bandwidth. The often quoted figure of more than 2 TB/s is memory bandwidth, and it belongs to the A100 80 GB variant. The A100 40 GB used in this scenario reaches about 1.6 TB/s of memory bandwidth. This is the preferred placement for communication-heavy tensor-parallel ranks because the GPUs can exchange tensors without using a normal data-center network path. However, putting two pods on the same node does not automatically prove that their GPUs share the expected NVLink domain. The machine type, GPU attachment, and physical topology still matter. Use `nvidia-smi topo -m` to inspect the topology visible inside the node. ### InfiniBand and RoCE InfiniBand is a high-throughput, low-latency network commonly used between GPU nodes. RDMA allows data to move between hosts with less CPU involvement and fewer copies. RoCE provides RDMA over a properly configured Ethernet fabric. These technologies can improve multi-node NCCL collectives dramatically compared with ordinary TCP networking. [NVIDIA Dynamo's multi-node guidance](https://docs.nvidia.com/dynamo/latest/kubernetes-deployment/scale/multinode-deployments) recommends InfiniBand, RoCE, or high-bandwidth Ethernet for strong performance. But InfiniBand is not magic. It normally exists inside a data center, availability zone, or dedicated cluster fabric. It does not create one lossless, microsecond-latency fabric between Iowa and Seoul. Even if both regions individually use InfiniBand, the path between those fabrics still crosses a wide-area network. Physical distance, routing, bandwidth limits, and congestion remain part of the equation. The practical rule is: > First place communicating GPUs inside the same accelerated network domain. Then use NVLink, NVSwitch, InfiniBand, or RoCE to make communication inside that domain faster. Fast networking without correct placement is wasted capability. Correct placement without a suitable network may still leave performance on the table. You need both. ## How KAI understands topology KAI builds a hierarchy from Kubernetes node labels. A simple cluster might look like this: ```text Region └── Zone └── Rack └── Host ``` You describe that hierarchy once with a cluster-scoped `Topology` custom resource: ```yaml apiVersion: kai.scheduler/v1alpha1 kind: Topology metadata: name: cluster-topology spec: levels: - nodeLabel: topology.kubernetes.io/region - nodeLabel: topology.kubernetes.io/zone - nodeLabel: cloud.example.com/topology-rack - nodeLabel: kubernetes.io/hostname ``` ![KAI Scheduler topology hierarchy built from Kubernetes node labels, with region, zone, rack, and host levels nested inside a cluster-scoped Topology resource](/images/kai-topology-hierarchy.webp "One Topology resource describes the region, zone, rack, and host levels | dark=/images/kai-topology-hierarchy-dark.webp") The order matters. The first label is the broadest domain, and every following level is nested inside the previous one. Every node considered at a required level must expose the relevant label. Check the node view before depending on it: ```bash kubectl get nodes -L topology.kubernetes.io/region \ -L topology.kubernetes.io/zone \ -L cloud.example.com/topology-rack \ -L kubernetes.io/hostname ``` A node without the required label can become ineligible rather than producing an obvious workload-level error. KAI's topology feature models domains expressed through node labels. With the classic `nvidia.com/gpu` device plugin, a host label does not describe the relationship between individual GPU devices. Newer integrations such as [NVIDIA Topograph](https://docs.nvidia.com/topograph/reference/node-labels-and-annotations) and DRA can expose accelerated interconnect domains, but the labels available depend on the platform. ## When should placement be required or preferred? KAI provides two complementary placement rules: ### Required placement Required placement is a hard boundary. All pods in the scheduling group must fit inside one domain at the selected level. ```yaml kai.scheduler/topology-required-placement: topology.kubernetes.io/region ``` If no region can fit the entire group, the group remains pending. KAI does not satisfy the job by placing half in Iowa and half in Seoul. ### Preferred placement Preferred placement is a soft locality goal. ```yaml kai.scheduler/topology-preferred-placement: topology.kubernetes.io/zone ``` KAI attempts to place the pods in one zone. If that is impossible, it can use another valid placement inside the required region instead of blocking the job. Together, the annotations read naturally: ```yaml metadata: annotations: kai.scheduler/topology: cluster-topology kai.scheduler/topology-required-placement: topology.kubernetes.io/region kai.scheduler/topology-preferred-placement: topology.kubernetes.io/zone ``` The job must remain in one region and should remain in one zone when capacity allows. The required rule does not select a specific region. If both Iowa and Seoul can fit the job, either may be valid. Add `nodeSelector` or `nodeAffinity` when the workload must run in a particular region. ## Why does distributed GPU training need gang scheduling? Topology rules are useful only when KAI understands which pods belong together. Suppose a two-pod job has a required region constraint, but each pod is scheduled independently. Pod A can satisfy the rule by running in Iowa. Pod B can also satisfy the same rule independently by running in Seoul. Both pods are individually inside one region, but the distributed job is still split. Gang scheduling turns those pods into one scheduling unit. KAI admits the required number of members together or keeps them pending together. ![Gang scheduling comparison: without a gang, two GPU pods each satisfy the region rule in different regions; with batch-min-member set to 2, KAI places both pods in one region or keeps them pending](/images/gang-scheduling-missing-piece.webp "Without a gang the pods can split across regions; batch-min-member keeps them together or pending | dark=/images/gang-scheduling-missing-piece-dark.webp") For a plain two-pod Job, set: ```yaml metadata: annotations: kai.scheduler/batch-min-member: "2" ``` KAI's [batch documentation](https://github.com/NVIDIA/KAI-Scheduler/blob/main/docs/batch/README.md) is explicit that a plain multi-pod Job is otherwise scheduled pod by pod, so this annotation is what turns the Job into one gang. The minimum member count must cover every pod that participates in the tightly coupled operation. A threshold of one cannot protect a two-worker collective. Framework-aware resources such as `PyTorchJob`, `JobSet`, or `LeaderWorkerSet` are usually better foundations for real distributed workloads because they also model worker roles and coordination. ## A complete two-GPU example ```yaml apiVersion: batch/v1 kind: Job metadata: name: two-gpu-training labels: kai.scheduler/queue: default-queue annotations: kai.scheduler/topology: cluster-topology kai.scheduler/topology-required-placement: topology.kubernetes.io/region kai.scheduler/topology-preferred-placement: topology.kubernetes.io/zone kai.scheduler/batch-min-member: "2" spec: parallelism: 2 completions: 2 completionMode: Indexed template: metadata: labels: kai.scheduler/queue: default-queue spec: schedulerName: kai-scheduler restartPolicy: Never containers: - name: trainer image: nvcr.io/nvidia/pytorch:24.12-py3 resources: limits: nvidia.com/gpu: 1 command: ["bash", "-lc", "python /workspace/train.py"] ``` This manifest establishes four important properties: 1. The pods are routed to `kai-scheduler`. 2. The workload is associated with a KAI queue. 3. Both workers form a gang with a minimum size of two. 4. The gang must fit inside one region and should fit inside one zone. The example focuses on scheduling. A production distributed training job still needs service discovery, rendezvous configuration, ranks, storage, credentials, and failure handling. ## How KAI chooses a domain Once the complete gang is known, KAI filters out topology domains that cannot satisfy it. Imagine the cluster has this free capacity: | Region | Free GPUs | Job request | Feasible? | | --- | ---: | ---: | --- | | Iowa | 2 | 2 | Yes | | Seoul | 2 | 2 | Yes | Both regions can fit the complete gang, so KAI may choose either. The guarantee is locality, not geography. The choice is not random, though. When more than one domain is feasible, KAI applies a bin-packing strategy: it prefers the domain with less free capacity relative to the request. That consolidates workloads into domains that are already relatively full and preserves emptier domains for future jobs. Within the selected domain, node ordering flips direction. Sub-domains at the preferred level with more free capacity are prioritized so that as many pods as possible share the same zone or rack. Now change the picture slightly. Suppose another workload lands first and takes one of the two GPUs in Iowa. When our two-GPU job arrives, the free capacity looks like this: | Region | Free GPUs | Job request | Feasible? | | --- | ---: | ---: | --- | | Iowa | 1 | 2 | No | | Seoul | 2 | 2 | Yes | Iowa now has only one free GPU, and the gang needs two together, so Iowa is no longer a valid domain. Seoul still has two free GPUs, so it can fit the whole gang. KAI places both workers in Seoul instead of combining the one free Iowa GPU with one free Seoul GPU. ![KAI Scheduler domain feasibility check: a two-GPU gang is rejected by a region with one free GPU and placed together in a region with two free GPUs](/images/kai-domain-feasibility.webp "KAI rejects a region that cannot fit the whole gang and places it where the gang fits | dark=/images/kai-domain-feasibility-dark.webp") That may temporarily leave a GPU unused in Iowa, but it avoids creating a distributed job whose communication path destroys its efficiency. Higher raw allocation is not the same as higher useful GPU utilization. ## Tighter placement for tensor parallel groups A region is only the outer boundary. Communication-heavy workloads often need tighter locality. For example: - Keep the whole workload inside one zone. - Keep each tensor-parallel group inside one rack or accelerated fabric domain. - Prefer NVLink-connected GPUs for ranks that exchange tensors on every layer. KAI supports [segment-based placement](https://github.com/NVIDIA/KAI-Scheduler/blob/main/docs/topology/segments.md) for symmetric groups. A 16-worker `PyTorchJob` with tensor-parallel groups of four can use segments: ```yaml apiVersion: kubeflow.org/v1 kind: PyTorchJob metadata: name: tp4-training annotations: kai.scheduler/topology: cluster-topology kai.scheduler/topology-required-placement: topology.kubernetes.io/zone spec: pytorchReplicaSpecs: Master: replicas: 1 template: metadata: labels: kai.scheduler/queue: default-queue spec: schedulerName: kai-scheduler containers: - name: pytorch image: pytorch/pytorch:latest Worker: replicas: 16 template: metadata: labels: kai.scheduler/queue: default-queue annotations: kai.scheduler/segment-size: "4" kai.scheduler/segment-topology-required-placement: cloud.example.com/topology-rack spec: schedulerName: kai-scheduler containers: - name: pytorch image: pytorch/pytorch:latest ``` The outer rule keeps the whole job inside one zone. The segment rule keeps each group of four workers inside one rack. Different groups may use different racks while every group remains intact. ![Segment-based placement: a 16-worker PyTorchJob inside one zone, with each tensor parallel group of four workers kept together on a single rack](/images/kai-segment-placement.webp "Each tensor parallel group of four stays on one rack inside the required zone | dark=/images/kai-segment-placement-dark.webp") A few practical rules apply. Segments are assigned by replica index, so `segment-size: "4"` puts workers 0 to 3 in the first segment, 4 to 7 in the second, and so on. The value must be greater than 1 and no larger than the replica count. Segmentation is currently wired up for the `Worker` replica of a `PyTorchJob` and for `LeaderWorkerSet`; other workload kinds need a hand-written `PodGroup` with explicit subgroups. A softer `kai.scheduler/segment-topology-preferred-placement` variant also exists for cases where rack locality is desirable but not mandatory. This matters because communication patterns are not uniform. Tensor-parallel ranks may communicate during every model layer, while data-parallel groups synchronize at different points. [Disaggregated inference](https://docs.nvidia.com/dynamo/latest/kubernetes-deployment/scale/topology-aware-scheduling) has another pattern: prefill and decode workers exchange KV-cache data, but communication inside each tensor-parallel group can be more frequent. Topology should follow the communication graph, not just the number of pods. ## Five checks before trusting the guarantee A required region rule prevents cross-region placement only when all of the following are true: 1. **KAI schedules the pods.** The pod spec contains `schedulerName: kai-scheduler`. 2. **The workload belongs to a KAI queue.** The expected queue label is present where the installed KAI version and workload integration require it. 3. **The communicating pods form one PodGroup.** KAI must treat them as one workload. 4. **The gang threshold covers the complete group.** A threshold of one cannot protect a two-worker collective. 5. **The topology exists in the scheduler's node view.** Required region, zone, rack, or fabric labels must be visible and correct. Also place topology annotations on the object level documented for the workload controller you use. Owner-based workloads may read constraints from their top-level owner rather than from individual pod templates. ## Benefits beyond avoiding a slow job Topology-aware scheduling improves more than raw training time: - **Higher useful GPU utilization:** GPUs spend more time computing and less time waiting for remote ranks. - **More predictable completion time:** Placement is constrained before the job starts instead of leaving performance to chance. - **Lower network cost:** Cross-region data transfer can add cloud egress or inter-region charges. - **Better failure behavior:** A complete gang waits when suitable capacity is unavailable instead of starting in a degraded topology. - **Improved cluster efficiency:** The scheduler preserves suitable domains for workloads that actually need them. - **Clearer platform contracts:** Teams can request locality as part of the workload rather than depending on tribal knowledge about node names. There is a trade-off. Stronger constraints can increase queue time because the scheduler may wait for a complete domain to become available. That is often the correct trade for synchronous distributed jobs. A job that starts immediately but runs several times slower does not represent useful scheduling efficiency. ## Final thoughts Distributed GPU performance comes down to placement: fast networking like NVLink or InfiniBand can only speed up communication once the workers that talk to each other actually sit in the same domain. KAI Scheduler fixes that at the source: the `Topology` resource describes the map, required and preferred placement set the boundary, gang scheduling applies it to the complete job, and segment placement keeps each communication group tight. The goal is not one node for everything, but each communicating group in the closest domain that matches its pattern. That pays off just as much for inference as for training. Disaggregated serving splits prefill and decode into separate pods that exchange KV-cache data over the network, so required placement keeps the serving group inside one region or zone, gang scheduling brings a replica up only when every role can be placed together, and segment placement keeps each tensor-parallel group on one rack or NVLink domain. The result is steadier tail latency and a lower cost per request.