Add configurable allocation policy (packed/distributed) for replicated and MIG resources

[wkd-woo]

Summary

• Add --allocation-policy flag (ALLOCATION_POLICY env) with distributed (default) and packed options
• packed mode bin-packs replicated/MIG devices onto fewest physical GPUs, freeing up remaining GPUs for full-GPU workloads
• Default behavior (distributed) is unchanged — no breaking changes

Motivation

The current distributedAlloc was designed for time-slicing, where distributing replicas across physical GPUs avoids compute contention. However, MIG devices also fall into this code path simply because AlignedAllocationSupported() returns false for them — not because distributed allocation is the right strategy.

MIG instances are hardware-isolated partitions with dedicated SMs and memory. Packing them onto fewer physical GPUs has no performance penalty, and frees up remaining GPUs for full-GPU workloads:

┌─────────────────────────────────────────────────────────────────────┐
│                Distributed Allocation (current default)              │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│  GPU 0            GPU 1            GPU 2            GPU 3           │
│  ┌──────────┐    ┌──────────┐    ┌──────────┐    ┌──────────┐      │
│  │ ████░░░░ │    │ ████░░░░ │    │ ████░░░░ │    │ ░░░░░░░░ │      │
│  │ MIG 1/5  │    │ MIG 1/5  │    │ MIG 1/5  │    │ MIG 0/5  │      │
│  └──────────┘    └──────────┘    └──────────┘    └──────────┘      │
│   ⚠ partial       ⚠ partial       ⚠ partial       ░ empty          │
│                                                                     │
│  → Full GPU request arrives: only 1 GPU available (GPU 3)           │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────┐
│                    Packed Allocation (bin-packing)                   │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│  GPU 0            GPU 1            GPU 2            GPU 3           │
│  ┌──────────┐    ┌──────────┐    ┌──────────┐    ┌──────────┐      │
│  │ ████████ │    │ ░░░░░░░░ │    │ ░░░░░░░░ │    │ ░░░░░░░░ │      │
│  │ MIG 3/5  │    │ MIG 0/5  │    │ MIG 0/5  │    │ MIG 0/5  │      │
│  └──────────┘    └──────────┘    └──────────┘    └──────────┘      │
│   ✓ packed         ★ free          ★ free          ★ free           │
│                                                                     │
│  → Full GPU request arrives: 3 GPUs available (GPU 1, 2, 3)        │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘

Relates to #491

Design

getPreferredAllocation
├─ Full GPU (AlignedAllocationSupported && no annotations)
│   └─ alignedAlloc (topology-based, unaffected by this change)
│
└─ MIG / time-slicing / MPS
├─ allocationPolicy=packed  → packedAlloc (new)
└─ allocationPolicy=distributed (default) → distributedAlloc (existing, unmodified)

Key decisions:

• distributedAlloc is completely untouched — packedAlloc is a separate function following the existing alignedAlloc/distributedAlloc pattern
• Full GPU nodes are unaffected — alignedAlloc is selected before allocationPolicy is ever checked, so setting packed on a full-GPU node has no effect
• Flag applies uniformly — when packed is set, it applies to MIG, time-slicing, and MPS. Silently ignoring a user-set flag for specific device types would be inconsistent
• Per-node config supported — works with existing config-manager + ConfigMap + node label (nvidia.com/device-plugin.config) mechanism via YAML config

Usage

CLI flag / Environment variable

--allocation-policy=packed
# or
ALLOCATION_POLICY=packed
Config file (per-node via ConfigMap + node label)

version: v1
flags:
  migStrategy: mixed
  plugin:
    allocationPolicy: packed

kubectl label node mig-node nvidia.com/device-plugin.config=mig-packed

Test plan

• TestDistributedAlloc (6 cases) — existing distributed behavior regression
• TestDistributedAllocIsDefault — default consistency verified over 10 iterations
• TestPackedAlloc (6 cases) — bin-packing: same GPU priority, overflow to next GPU
• TestPackedVsDistributedContrast — two strategies produce different results on same input
• TestFullGPUNodeIgnoresAllocationPolicy (3 sub-cases) — full GPU / MIG / replicated branch path verification
  All existing internal/rm/ tests pass unchanged

---

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

[ottowhite]

Hi! Very cool idea, I actually really want a very similar custom allocation policy for replicated (non-MIG) resources, but for enabling bin-packing replicas with topology-awareness.

For example, this would could enable the allocation of 1.6 GPUs to a single model, with two tensor parallel ranks across 2 GPUs, 0.8 GPUs allocated to each tensor parallel tank, and both within a single NVLink domain (see my pretty diagram for an example). This would avoid excessive communication overhead during matrix multiplications across both tensor parallel model ranks, while leaving the 0.4 GPUs left over to be allocated to a smaller model, decreasing latency and increasing throughput in (surprisingly common) multi-model deployment scenarios including RAG with LLMs + Embedding models.

I believe the implementation of this policy is essentially a composite between what you've implemented, and the alignedAlloc code path which uses go-gpuallocator (gaining topology-awareness through NVML). At present, any of the replicated resource code paths discard topology awareness in favour of evenly load balancing GPU fractions with the distributed allocation policy, which prevents the use-case that I'm after.

Would be great if you're interested in discussing this more. I wonder if it's possible to slip it in with this PR? Also happy to contribute if that helps.

[ottowhite]

Was the only change in this function and distributedAlloc the comparator between idiff and jdiff? Can we just pass in an enum and then have a condition on the enum to dictate this behavior? There is a lot of code duplication at the moment.

[wkd-woo]

Thanks for the review!

The duplication is intentional — I wanted distributedAlloc to stay completely untouched so reviewers can verify there's zero behavioral change just by looking at the diff.

On merging them: the two functions share the same structure because packed is just the inverse of distributed (flipped comparator). But as you mentioned with topology-aware bin-packing, a future policy would need NVML topology data and different sorting logic entirely — so I'd rather keep them separate now than merge and split again later.

That said, happy to extract the common setup (candidate filtering + replica counting) into a shared helper if that feels cleaner.

[ottowhite]

No problem :) to be honest I think it would actually be easier to review/maintain if you generalised the allocation function slightly.

It could be generalised to fractionalAlloc, taking an enum that allows distributedFractionPlacement or packedFractionPlacement policies. Then internally you just pick the comparator based on the policy.

What you suggested with pulling out the common helpers such as candidate filtering and replica counting is probably even better than the generalisation I suggested, as it enables the policies to be kept completely separate, and is more extensible in that it enables easier implementation of new policies again such as the topology-aware one.

When I was reviewing the code, I pulled up both distributedAlloc and packedAlloc side by side and was cross-checking for a while before I saw the flipped comparator. I expect other reviewers and contributors to go through a similar process. This is just my two cents though.

You've also improved the state of testing so that's an additional signal that this is functionally equivalent to what there was before.

In terms of the different topology-aware placement policy, I think it makes sense to keep it separate for now. But cool that you're opening this up and making it more configurable! I may build off of this myself.

[wkd-woo]

@ottowhite I've extracted prepareCandidates in the second commit.
Each policy function now only has its sorting logic, so the difference should be immediately visible

[ottowhite]

Looking better! There's still quite a bit of duplication across both of the Alloc functions. What is common is that they're both greedy allocation algorithms (repeatedly evaluating and selecting the next best option, without looking further ahead). Could you extract this into another helper greedyDeviceAlloc or something similar that makes the difference between these two functions even clearer? Could pass in a candidate comparator. It will also be much clearer to extend with other greedy allocation policies. For example the one that I spoke about is another greedy allocation algorithm that happens to use topology-awareness in it's comparator.

[rajatchopra]

I agree that we should have a map of comparator functions to choose from. And a helper function that keeps the common portion in one place.

[wkd-woo]

Done in 4afcd2e :) distributedAlloc/packedAlloc are gone — there is now a single greedyAlloc that takes a comparator, plus an allocationComparators map keyed by policy. The difference between the two policies ends up being just < vs > in the map. Unknown policies fall back to distributed (covered by a new test), and the existing tests pass unchanged.

[ottowhite]

Also, is there a reason why the alignedAlloc happens to be in nvml_manager.go? It seems like just another allocation policy that would be more suited to be in allocate.go rather than there. I know this is not due to your code change but moving it there could improve the organisation of the code, and ease to extend with different allocation policies.

[wkd-woo]

Thank @ottowhite for your review :)
On moving alignedAlloc — agreed it could improve organization, but I'd rather keep this PR focused on the new feature and not restructure existing code!

I'd like to hear from a maintainer at this point to see if the current direction works.

[ottowhite]

Fair enough! Though I think those refactors would be a minor improvement for organisation/extensibility, I think the testing is strong and implementation is solid.

LGTM!

Will be great to hear what the maintainers have to say. Also really keen on this feature and the more extensible policy opportunities that it opens up.

[ottowhite]

Hey! By the way we're running your packedAlloc allocation policy and it's working for multi-GPU deployments. We will imminently be forking and building off it it.

[wkd-woo]

Hi @elezar, just checking in to see if there's anything else needed from my side to move this forward. I'm happy to address any feedback :)

[wkd-woo]

Hi @rajatchopra, could you take a look at this PR when you have a chance?
I'd appreciate your feedback.

[rajatchopra]

/lgtm

Pending:

• rebase after #1788 merges
• squash commits

[tariq1890]

@wkd-woo Please rebase your PR on top of the latest main.

[wkd-woo]

@rajatchopra @tariq1890 Rebased onto latest main (now includes #1788) and squashed into a single commit as requested.

Heads-up — this was not a clean rebase, so this needs another look before merge. #1788 landed the day after your /lgtm and it touches the exact same allocation path this PR refactors, so I had to reconcile them by hand:

• fix distributedAlloc tie-break to preserve distribution across physical GPUs #1788 added a pickedFrom tie-break to distributed allocation (rotate to a sibling physical GPU when two GPUs tie on allocated-replica count) plus a regression test TestDistributedAlloc_PartiallyAllocated_DistributesAcrossDistinctGPUs.
• This PR had already replaced distributedAlloc/packedAlloc with a single greedyAlloc(comparator), so fix distributedAlloc tie-break to preserve distribution across physical GPUs #1788's fix no longer had a place to live.

To keep both, I folded the tie-break into greedyAlloc as a secondary sort key applied only when the comparator ranks two GPUs equally:

 sort.Slice(candidates, func(i, j int) bool {
     iid := AnnotatedID(candidates[i]).GetID()
     jid := AnnotatedID(candidates[j]).GetID()
-    return preferred(replicas[iid], replicas[jid])
+    ri, rj := replicas[iid], replicas[jid]
+    if preferred(ri, rj) {
+        return true
+    }
+    if preferred(rj, ri) {
+        return false
+    }
+    // comparator ranks the two GPUs equally -> prefer the physical
+    // device touched least during this allocation
+    return pickedFrom[iid] < pickedFrom[jid]
 })

This is the only logic change beyond the code you approved. It's required rather than optional: without it, #1788's regression test now fails on this branch (both replicas stack onto GPU-0); with it, the distributed policy keeps spreading. It's a no-op for packed, whose comparator dominates once packing starts. All internal/rm tests pass (TestDistributedAlloc, TestPackedAlloc, TestDistributedAllocIsDefault, TestPackedVsDistributedContrast, TestFullGPUNodeIgnoresAllocationPolicy, and #1788's TestDistributedAlloc_PartiallyAllocated_...).

PTAL when you get a chance.

[wkd-woo]

cc @tariq1890 — flagging for the merge: this was rebased onto main after #1788 with one small logic change beyond @rajatchopra's /lgtm (folded #1788's distributed-allocation tie-break into this PR's new greedyAlloc — details and diff in the comment above). It's a no-op for packed, and all internal/rm tests pass including #1788's regression test. Should be ready to merge once CI is green — thanks!

[rajatchopra]

@wkd-woo >> To keep both, I folded the tie-break into greedyAlloc as a secondary sort key applied only when the comparator ranks two GPUs equally
That is good work resolving the conflicts. Thanks. Needs another benign rebase though.

[rajatchopra]

/lgtm

[wkd-woo]

@rajatchopra Rebased onto latest main. Clean, no conflicts, and all tests pass. Thanks!

cc. @tariq1890

[tariq1890]

Suggested change

	return fmt.Errorf("invalid --allocation-policy option: %v", *config.Flags.Plugin.AllocationPolicy)
	return fmt.Errorf("invalid --allocation-policy option: %s", *config.Flags.Plugin.AllocationPolicy)

[wkd-woo]

👍
Will Apply.

[tariq1890]

Suggested change

	total, available int
	total int
	available int

[wkd-woo]

👍
Will Apply.

[tariq1890]

Can we come up with a more specific name here? The name ALLOCATION_POLICY does not lend itself to an allocation policy used for replicated and MIG devices.

[tariq1890]

How about SHARED_DEVICES_ALLOCATION_POLICY?

[wkd-woo]

Agreed.

Will rename to SHARED_DEVICES_ALLOCATION_POLICY / --shared-devices-allocation-policy, and update the config field to match.

[tariq1890]

Thank you !

[tariq1890]

Would this work?

Suggested change

	if preferred(ri, rj) {
	return true
	}
	if preferred(rj, ri) {
	return false
	if ri.allocated() != rj.allocated() {
	return preferred(ri, rj)
	}

[wkd-woo]

Yes, that works. Both comparators only order by allocated(), so it comes out the same. I went with the two-way check to avoid assuming what the comparator compares, but with just these two policies that's probably overkill.
Will apply it, and add a small test on the comparators to cover that assumption.

[tariq1890]

@wkd-woo Thank you very much for your patience with this PR and I appreciate all the work you've put in. Just a few more comments, but we are definitely looking to merge this and ensure that this lands in the next release of the device-plugin.

[wkd-woo]

@tariq1890 Applied all four in b8c575ef, tests pass 🙏

[tariq1890]

Can you squash your commit history?

[wkd-woo]

@tariq1890 Squashed !

[tariq1890]

/ok to test 196c7c0

[tariq1890]

@wkd-woo. Your commit is not signed. See here for context.

[wkd-woo]

@tariq1890 Signed the commit. Thank you for your patience.

[tariq1890]

/ok to test 891d77f