feat: SM Limiting using CUDA Green Contexts

[anandj91]

Overview

CUDA 13 introduced green contexts (cuGreenCtxCreate / cuCtxFromGreenCtx), which are a lightweight alternative to traditional contexts, that can be used to select a subset of device resources. This allows the developer to, for example, select SMs from distinct spatial partitions of the GPU and target them via CUDA stream operations, kernel launches, etc. As opposed to the current token-based rate limiter backed by an NVML-polling watcher thread in HAMi-core, the SM allocation in green contexts is driver enforced with no polling or per-launch coordination overhead.

Proposed Design

A new USE_GREEN_CTX=1 environment variable gates the feature. When set, the cuCtxCreate hook (already in place) will compute target_sms = rount(total_sms * CUDA_DEVICE_SM_LIMIT / 100), create a green context scoped to that count, and return the derived CUcontext to the caller. The existing CUDA_DEVICE_SM_LIMIT environment variable is reused as-is. If cuGreenCtxCreate fails (e.g., on a driver older than CUDA 13.1), the hook returns the error explicitly rather than silently falling back. Default behavior is unchanged. Without USE_GREEN_CTX=1, everything works exactly as before.

[anandj91]

Raised this PR as a precursor to this feature proposal: #178

[anandj91]

@archlitchi would like to get some thoughts on this feature request. if what I propose sounds good, I can raise a PR for the same. Wanted to make sure that we are aligned on this.

[archlitchi]

I've tried cuGreenContext about two months ago, and the result is:

it can only be shared among multiple contexts in the same process

For example, Process A uses green context with 1/2 of total SMs. And yes, it can be executed simultaneously with another green context which allocates another 1/2 of total SMs in Process A. But it can't be run simultaneously with green context in process B.

So it can't replace the token-limiter right now. But i won't deny if this problem solves, it's a far better design than the original, perhaps you could have a try and let me know the result:)

[anandj91]

@archlitchi we have a paper published on this topic: https://dl.acm.org/doi/10.1145/3669940.3707282

In this work, what we have done is to maintain the shared cudaContext for all the processes using a GPU. Then we implement our own rate limiting strategy by performing live kernel transformations as described in the paper. I believe we can simply replace kernel transformation approach using the cudaGreenContext.

One downside of this approach is, sharing cuda context means there is no isolation between processes using the GPU. If one of the processes crashes, it can crash others too.

[archlitchi]

https://dl.acm.org/doi/10.1145/3669940.3707282

yes, it looks like mps, have you compared to them?

[archlitchi]

@archlitchi we have a paper published on this topic: https://dl.acm.org/doi/10.1145/3669940.3707282

In this work, what we have done is to maintain the shared cudaContext for all the processes using a GPU. Then we implement our own rate limiting strategy by performing live kernel transformations as described in the paper. I believe we can simply replace kernel transformation approach using the cudaGreenContext.

One downside of this approach is, sharing cuda context means there is no isolation between processes using the GPU. If one of the processes crashes, it can crash others too.

And i'm thinking if it's possible to implement a remote-cuda by using your sharedContext? We are having a discussion about it, maybe you're interested? https://gist.github.com/ManishSharma1609/92775165012b9d67e2ea707d9c7aa44a