Add spread allocation policy for maximum per-pod distinct-GPU coverage

[jonathan-meiri]

Summary

• Add spread as a third value for --shared-devices-allocation-policy, alongside distributed (default) and packed.
• spread prefers physical GPUs the current pod has not yet picked from, maximizing distinct-GPU coverage for a single multi-slot allocation.
• Default behavior (distributed) is unchanged.

Depends on #1826 (heap refactor) — the current diff includes that commit until #1826 lands, then collapses to just the spread policy addition. Opened as draft for now.

Contributed by @Meiri28 on behalf of @runatom-ai.

Motivation

distributed minimizes cluster-wide load imbalance — it prefers the GPU with the fewest replicas already allocated across all pods. packed does the opposite for bin-packing. Neither guarantees that a single pod's multi-slot request touches distinct physical GPUs.

Consider a node where GPU-0 has 3 free replicas (5 allocated) and GPU-1 has 5 free (3 allocated). A pod requests 2:

┌─────────────────────────────────────────────────────────────────────┐
│              distributed (current default)                          │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│  GPU 0                       GPU 1                                  │
│  ┌──────────┐               ┌──────────┐                            │
│  │ █████░░░ │               │ ███░░░░░ │      before                │
│  │  5 alloc │               │  3 alloc │                            │
│  └──────────┘               └──────────┘                            │
│                                                                     │
│  Pod requests 2 slots → distributed picks GPU with fewer allocated  │
│                                                                     │
│  ┌──────────┐               ┌──────────┐                            │
│  │ █████░░░ │               │ █████░░░ │      after                 │
│  │  5 alloc │               │  5 alloc │                            │
│  └──────────┘               └──────────┘                            │
│   ⚠ untouched                 ✓ 2 slots (pod concentrated)          │
│                                                                     │
│  → Pod runs entirely on GPU 1. No NVLink/NCCL across GPUs.          │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────┐
│              spread (new)                                           │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│  GPU 0                       GPU 1                                  │
│  ┌──────────┐               ┌──────────┐                            │
│  │ █████░░░ │               │ ███░░░░░ │      before                │
│  │  5 alloc │               │  3 alloc │                            │
│  └──────────┘               └──────────┘                            │
│                                                                     │
│  Pod requests 2 slots → spread prefers GPUs untouched by THIS pod   │
│                                                                     │
│  ┌──────────┐               ┌──────────┐                            │
│  │ ██████░░ │               │ ████░░░░ │      after                 │
│  │  6 alloc │               │  4 alloc │                            │
│  └──────────┘               └──────────┘                            │
│   ★ 1 slot                    ★ 1 slot (pod spans both GPUs)        │
│                                                                     │
│  → Multi-GPU workloads (NCCL, DDP, tensor-parallel) get real        │
│    hardware diversity for a single request.                         │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘

Concrete workloads that benefit:

• Data-parallel / distributed training — each rank on a distinct physical GPU so ring-allreduce / NCCL spans real hardware
• Fault-isolated inference — a pod requesting 2 GPUs specifically to survive a single-card failure
• NVLink topology plays — a first step toward topology-aware allocation (weak preference: at least distinct cards)

@rajatchopra explicitly floated additional policies on #1788:

"we can decide on distributed allocation or a packed one, or a topology driven one"

This adds the third policy that #1621's abstraction was built for.

Design

Small refactor: the replicaComparator signature is enriched from func(i, j *replicaCount) bool to func(i, j *gpuAllocState) bool. Each policy now owns both its primary ordering and tie-break — the queue's Less is a pure passthrough:

distributed: func(i, j *gpuAllocState) bool {
    if i.count.allocated() != j.count.allocated() {
        return i.count.allocated() < j.count.allocated()   // primary
    }
    return i.pickedFrom < j.pickedFrom                     // tie-break
}
packed: func(i, j *gpuAllocState) bool {
    if i.count.allocated() != j.count.allocated() {
        return i.count.allocated() > j.count.allocated()   // primary
    }
    return i.pickedFrom < j.pickedFrom                     // tie-break
}
spread: func(i, j *gpuAllocState) bool {
    if i.pickedFrom != j.pickedFrom {
        return i.pickedFrom < j.pickedFrom                 // primary
    }
    return i.count.allocated() < j.count.allocated()       // tie-break
}

distributed and packed are behavior-preserving — the primary+tie-break they used to get from greedyAlloc's wrapping is now spelled out in the comparator body itself. All existing tests (TestDistributedAlloc, TestPackedAlloc, TestPackedVsDistributedContrast) pass unchanged.

Usage

--shared-devices-allocation-policy=spread
# or
SHARED_DEVICES_ALLOCATION_POLICY=spread

Config file:

version: v1
flags:
  plugin:
    sharedDevicesAllocationPolicy: spread

Commits are DCO-signed.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[tariq1890]

THanks @jonathan-meiri ! Can you rebase the PR?

[jonathan-meiri]

Rebased on latest main 🙏
Thanks for the review, @tariq1890!

[abrarshivani]

@jonathan-meiri Thanks for working on this. LGTM.

[abrarshivani]

/ok to test 046b76b

needs more work

[tariq1890]

What about MIG devices?

With MIG, spread would treat each MIG UUID as a distinct physical GPU, yes?. Candidate buckets use device UUIDs, while multiple MIG devices can share one parent. A multi-slot allocation can remain on one physical GPU.

[jonathan-meiri]

Good catch, you're right. Buckets key on the device UUID (GetID), and each MIG instance has its own UUID, so spread spans distinct MIG instances that can share one physical GPU. This is pre-existing (distributed/packed too).

I've kept this PR to the spread comparator and opened #2036 to bucket by parent GPU (from Device.Index) for all policies.

Let me know if you'd prefer to fold that into this PR, or if you're happy keeping it separate.

[tariq1890]

If I am not wrong, this changes the tie-breaking logic and it would apply to the distributed and packed allocation policies. This goes beyond the scope of this PR - which is an introduction of a new allocation policy

[jonathan-meiri]

Good point, seeding pickedFrom leaked into distributed/packed. Fixed: pickedFrom stays pure, required accounting is now spread-only (touched()), with a test to lock it in.

Packed's tie-break could improve too, but that's pre-existing. I'll follow up separately.

[tariq1890]

This godoc needs to be updated now that the spread comparator has been updated to use the touched() method as a primary sort key.

[tariq1890]

/ok to test 7062bd0

[jonathan-meiri]

Thanks @tariq1890 and @abrarshivani for the thorough reviews and the merge, really appreciate it!

The MIG parent-GPU bucketing follow-up we discussed is up as #2036 whenever you have a chance.