Topology-Aware GPU scheduling with KAI Scheduler

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:

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 and gang scheduling. 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.

For the real value in your environment, measure the VM-to-VM path. Google Cloud's Performance Dashboard 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:

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:

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:

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, 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 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:

Region
└── Zone
    └── Rack
        └── Host

You describe that hierarchy once with a cluster-scoped Topology custom resource:

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

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:

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 and DRA can expose accelerated interconnect domains, but the labels available depend on the platform.

#Required and preferred placement

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.

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.

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:

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.

#The missing piece: 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.

For a plain two-pod Job, set:

metadata:
  annotations:
    kai.scheduler/batch-min-member: "2"

KAI's batch documentation 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

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:

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:

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.

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 for symmetric groups. A 16-worker PyTorchJob with tensor-parallel groups of four can use segments:

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.

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 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.