Add configurable allocation policy (packed/distributed) for replicated and MIG resources by wkd-woo · Pull Request #1621 · NVIDIA/k8s-device-plugin

wkd-woo · 2026-02-11T09:05:26Z

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 · 2026-02-11T09:05:30Z

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 · 2026-02-15T02:15:01Z

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 · 2026-02-16T11:21:52Z

+			jid := AnnotatedID(candidates[j]).GetID()
+			idiff := replicas[iid].total - replicas[iid].available
+			jdiff := replicas[jid].total - replicas[jid].available
+			return idiff > jdiff


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.

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.

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.

@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

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.

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.

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 · 2026-02-17T13:32:13Z

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 · 2026-02-17T13:50:53Z

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 · 2026-02-17T14:10:03Z

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 · 2026-02-19T18:26:24Z

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 · 2026-02-25T14:25:34Z

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 · 2026-06-16T03:05:43Z

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 · 2026-06-30T23:10:36Z

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

wkd-woo · 2026-07-01T03:12:57Z

@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 · 2026-07-01T03:16:41Z

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 · 2026-07-20T06:37:01Z

@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 · 2026-07-20T08:12:12Z

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

cc. @tariq1890

tariq1890 · 2026-07-20T23:26:44Z

+		case spec.AllocationPolicyDistributed:
+		case spec.AllocationPolicyPacked:
+		default:
+			return fmt.Errorf("invalid --allocation-policy option: %v", *config.Flags.Plugin.AllocationPolicy)


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)

👍
Will Apply.

tariq1890 · 2026-07-20T23:32:28Z

-	// Get the set of candidate devices as the difference between available and required.
+// replicaCount tracks the total and available replica counts for a physical GPU.
+type replicaCount struct {
+	total, available int


Suggested change

total, available int

total int

available int

👍
Will Apply.

tariq1890 · 2026-07-20T23:56:58Z

+			Name:    "allocation-policy",
+			Value:   spec.AllocationPolicyDistributed,
+			Usage:   "the allocation policy for replicated and MIG resources:\n\t\t[distributed | packed]",
+			EnvVars: []string{"ALLOCATION_POLICY"},


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.

How about SHARED_DEVICES_ALLOCATION_POLICY?

Agreed.

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

Thank you !

tariq1890 · 2026-07-21T00:08:31Z

+			if preferred(ri, rj) {
+				return true
+			}
+			if preferred(rj, ri) {
+				return false


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)

}

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 · 2026-07-21T00:10:51Z

@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 · 2026-07-21T02:43:33Z

@tariq1890 Applied all four in b8c575ef, tests pass 🙏

tariq1890 · 2026-07-21T05:19:44Z

Can you squash your commit history?

wkd-woo · 2026-07-21T05:26:11Z

@tariq1890 Squashed !

tariq1890 · 2026-07-21T15:27:02Z

/ok to test 196c7c0

tariq1890 · 2026-07-21T16:37:12Z

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

…d and MIG resources The existing allocation spreads replicated devices evenly across physical GPUs (distributed), which was designed for time-slicing where workloads compete for shared compute. MIG instances, however, are hardware-isolated and do not suffer from contention when packed onto the same GPU. This adds a --shared-devices-allocation-policy flag (env: SHARED_DEVICES_ALLOCATION_POLICY) with two options: "distributed" (default, preserving current behavior) and "packed" (bin-packing onto the fewest physical GPUs). The packed policy frees up entire GPUs for full-GPU workloads in mixed clusters. The flag applies uniformly to all non-aligned allocation paths (MIG, time-slicing, MPS) and can be configured per-node via ConfigMap and the nvidia.com/device-plugin.config node label. Full-GPU nodes are unaffected: alignedAlloc is selected before the policy is ever consulted. Both policies share a single greedyAlloc helper that differs only in a replicaComparator, selected from an allocationComparators map keyed by policy (comparatorForPolicy falls back to distributed for unknown values). This keeps the strategies expressed as a flipped comparator and lets future policies be added as new map entries without duplicating the selection loop. When the comparator ranks two physical GPUs equally, a per-allocation pickedFrom tie-break rotates to the least-touched sibling, preserving distribution across physical GPUs. Relates to NVIDIA#491 Signed-off-by: wkd-woo <wkdwoos@gmail.com>

wkd-woo · 2026-07-21T17:18:46Z

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

tariq1890 · 2026-07-21T18:47:01Z

/ok to test 891d77f

Follow-up on top of NVIDIA#1621, which introduced the shared greedyAlloc loop with a pluggable replicaComparator (distributed vs packed). The loop still sorts the full candidate slice inside the allocation loop, paying O(n log n) per iteration for n iterations and giving O(n² log n) overall. Since all annotated replicas from the same underlying physical device share the same sort key, sorting at the replica granularity is wasted work — only m (the number of distinct physical devices contributing candidates) needs to be reordered. Refactor greedyAlloc to bucket candidates by their underlying physical device into a small gpuAllocState per device, holding a shared *replicaCount, the pickedFrom counter, and the remaining candidate IDs. A gpuPriorityQueue defers to the caller-supplied replicaComparator on allocated() for primary ordering and to pickedFrom for the tie-break (unchanged semantics). Each iteration pops the best device, takes one of its remaining replicas, updates counters, and pushes it back if any remain. Total cost drops to O(n log m). Both allocation policies (distributed and packed) benefit; no behavior change — the existing test suite (TestDistributedAlloc, TestPackedAlloc, TestPackedVsDistributedContrast, TestDistributedAlloc_PartiallyAllocated_DistributesAcrossDistinctGPUs, etc.) passes unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-Authored-By: runatom-ai <258621014+runatom-ai@users.noreply.github.com> Signed-off-by: Jonathan Meiri <33288957+Meiri28@users.noreply.github.com>

wkd-woo force-pushed the feature/packed-allocation-policy branch from 421e3c9 to 6686d1a Compare February 11, 2026 09:06

ottowhite reviewed Feb 16, 2026

View reviewed changes

wkd-woo force-pushed the feature/packed-allocation-policy branch 2 times, most recently from 8656ef3 to 2839bfc Compare February 17, 2026 12:22

jonathan-meiri mentioned this pull request May 19, 2026

Support per-physical-GPU replica counts in sharing.timeSlicing and sharing.mps #1786

Open

rajatchopra mentioned this pull request Jun 1, 2026

fix distributedAlloc tie-break to preserve distribution across physical GPUs #1788

Merged

wkd-woo requested a review from rajatchopra June 16, 2026 03:03

rajatchopra approved these changes Jun 29, 2026

View reviewed changes

wkd-woo force-pushed the feature/packed-allocation-policy branch from 4afcd2e to 59665d8 Compare July 1, 2026 03:12

rajatchopra mentioned this pull request Jul 20, 2026

perf(rm): replace per-iteration sort in distributedAlloc with a min-heap #1826

Open

rajatchopra approved these changes Jul 20, 2026

View reviewed changes

wkd-woo force-pushed the feature/packed-allocation-policy branch from 59665d8 to b2c3adf Compare July 20, 2026 08:09

tariq1890 reviewed Jul 21, 2026

View reviewed changes

wkd-woo force-pushed the feature/packed-allocation-policy branch from b2c3adf to b8c575e Compare July 21, 2026 02:41

wkd-woo force-pushed the feature/packed-allocation-policy branch from 14e9131 to 196c7c0 Compare July 21, 2026 05:25

tariq1890 approved these changes Jul 21, 2026

View reviewed changes

tariq1890 added this to the v0.20.0 milestone Jul 21, 2026

wkd-woo force-pushed the feature/packed-allocation-policy branch from 196c7c0 to 891d77f Compare July 21, 2026 17:13

tariq1890 merged commit 88a79d7 into NVIDIA:main Jul 21, 2026
38 of 41 checks passed

jonathan-meiri mentioned this pull request Jul 22, 2026

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

Draft

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

Uh oh!

Conversation

wkd-woo commented Feb 11, 2026

Summary

Motivation

Design

Usage

CLI flag / Environment variable

Uh oh!

copy-pr-bot Bot commented Feb 11, 2026

Uh oh!

ottowhite commented Feb 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ottowhite Feb 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ottowhite commented Feb 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wkd-woo commented Feb 17, 2026

Uh oh!

ottowhite commented Feb 17, 2026

Uh oh!

ottowhite commented Feb 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wkd-woo commented Feb 25, 2026

Uh oh!

wkd-woo commented Jun 16, 2026

Uh oh!

rajatchopra left a comment

Choose a reason for hiding this comment

Uh oh!

tariq1890 commented Jun 30, 2026

Uh oh!

wkd-woo commented Jul 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wkd-woo commented Jul 1, 2026

Uh oh!

rajatchopra commented Jul 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rajatchopra left a comment

Choose a reason for hiding this comment

Uh oh!

wkd-woo commented Jul 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

ottowhite commented Feb 15, 2026 •

edited

Loading

ottowhite Feb 17, 2026 •

edited

Loading

ottowhite commented Feb 17, 2026 •

edited

Loading

ottowhite commented Feb 19, 2026 •

edited

Loading

wkd-woo commented Jul 1, 2026 •

edited

Loading

rajatchopra commented Jul 20, 2026 •

edited

Loading

wkd-woo commented Jul 20, 2026 •

edited

Loading

tariq1890 commented Jul 21, 2026 •

edited

Loading

wkd-woo commented Jul 21, 2026 •

edited

Loading