Making Shared GPUs Even Safer with Kubex and HAMi-core

A few months ago, we introduced Kubex support for the KAI Scheduler to improve GPU sharing for production inference workloads.  

The basic model is simple:   

The KAI Scheduler handles placement and GPU sharing. Kubex continuously observes usage and adjusts those allocations as demand changes.   

KAI provides the scheduling foundation. It lets multiple workloads share a GPU while accounting for the amount of GPU each workload requests.  

Kubex then closes the loop. Instead of treating that GPU fraction as a static configuration, Kubex monitors actual compute and memory usage, rightsizes allocations, and rebalances workloads as demand changes.  

However, the system still had an important limitation: GPU memory allocations could be adjusted in real time, but those allocations were not enforced memory boundaries. 

A workload could request a fraction of a GPU, but CUDA did not prevent it from allocating more GPU memory than that fraction represented. KAI could decide how to share the GPU, but the workload could still access the underlying GPU memory. When a workload needed more memory, Kubex could respond by increasing its allocation, but there was no mechanism to prevent it from exceeding that allocation in the first place. 

Extending Kubex GPU sharing with HAMi-core   

KAI Scheduler now integrates with HAMi-core to enforce GPU memory limits for shared workloads. We have updated Kubex to support the latest KAI stack, allowing Kubex-managed GPU fractions to benefit from HAMi-core memory enforcement when it is enabled. 

HAMi-core is transparently injected into application containers at runtime; no application changes or recompilation are required. It intercepts CUDA memory allocations and enforces the memory limit associated with the GPU fraction continuously optimized by Kubex and scheduled by KAI. A workload that receives 4 GiB of GPU memory can no longer consume 8 GiB simply because the physical GPU has additional memory available.  

Previously, we had to rely on workloads to stay within their allocations or on Kubex to detect excess demand and increase those allocations. Now, HAMi-core enforces a memory boundary for application containers that share the same GPU.  

The stack now looks roughly like this:   

KAI Scheduler: decides where shared GPU workloads run and how much GPU capacity they receive.  

HAMi-core: enforces the GPU memory boundary inside each application container.  

Kubex: observes actual usage and continuously rightsizes and rebalances those allocations.  

This setup does not turn time-slicing into hardware partitioning like MIG. HAMi-core alone cannot stop workloads from sharing the underlying GPU or affecting each other through compute contention. Kubex helps address compute contention through continuous monitoring and rebalancing as demand changes.  

Memory remains one of the biggest risks when we increase GPU density, but HAMi-core now provides an enforced memory limit.  

Together, KAI Scheduler, HAMi-core, and Kubex provide a stronger foundation for running shared GPU inference in production on hardware you already have.  

If you want to learn more about the specifics of this integration, see our official documentation.

Get in touch with us to discuss your GPU use case and see how Kubex can help: kubex.ai/product/demo/