[CUDA] Add experimental green context support for SM carveout #159104
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/159104
Note: Links to docs will display an error until the docs builds have been completed. ❌ 8 New FailuresAs of commit 6db03ba with merge base 5f5f508 ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
test/test_matmul_cuda.py
Outdated
| @@ -1501,6 +1502,27 @@ def test_honor_sm_carveout(self) -> None: | |||
| self.assertNotEqual(no_carveout, carveout_66) | |||
| self.assertNotEqual(carveout_66, carveout_0) | |||
|
|
|||
| @unittest.skipIf(not _get_torch_cuda_version() >= (12, 8), "Green Context only tested on 12.8+") | |||
| def test_greencontext_caveout(self): | |||
There was a problem hiding this comment.
| def test_greencontext_caveout(self): | |
| def test_greencontext_carveout(self): |
test/test_matmul_cuda.py
Outdated
| @unittest.skipIf(not _get_torch_cuda_version() >= (12, 8), "Green Context only tested on 12.8+") | ||
| def test_greencontext_caveout(self): | ||
| a = torch.randn(4096, 4096, device='cuda', dtype=torch.bfloat16) | ||
| ctx = torch.cuda.green_contexts.GreenContext.create(0, 1) |
There was a problem hiding this comment.
want to make it _green_contexts or are we shore enough?
There was a problem hiding this comment.
One thing I could think of, num_sms should be the first arg, with the second defaulting to current device. Also, maybe use kwargs in the test to make it more obvious
There was a problem hiding this comment.
Not sure if this is the final API that we really want to expose over e.g., a context-manager API so users are not manually creating/destroying contexts or calling pybind'ed methods directly
| #if defined(CUDA_VERSION) && CUDA_VERSION >= 12080 | ||
| CUdevice device; | ||
| C10_CUDA_DRIVER_CHECK( | ||
| c10::cuda::DriverAPI::get()->cuDeviceGet_(&device, device_id)); |
There was a problem hiding this comment.
would this work if someone tries to create GreenContext before initializing cuda context
There was a problem hiding this comment.
nope, that was indeed broken
| c10::cuda::DriverAPI::get()->cuCtxPushCurrent_(context_)); | ||
| } | ||
| // currently hardcoes the new green context to use the default stream | ||
| // TODO(eqy): consider creating a new stream if e.g., it allows interop |
There was a problem hiding this comment.
anything preventing us from using a different stream now?
There was a problem hiding this comment.
Not that I'm aware of, but I think it gets messy as say the user is already in a stream context manager---in that case the new context implicitly switches to another "side" stream (even if it's the default stream from the perspective of the new context). We'll still sync appropriately before and after switching to the green context, but if they are checking things like torch.cuda.current_stream it may cause some confusion
There was a problem hiding this comment.
but it would still be as confusing to get 0 for torch.cuda.current_stream() which I think would happen now?
torch/csrc/cuda/green_context.h
Outdated
|
|
||
| // Implement move operations | ||
| GreenContext(GreenContext&& other) noexcept | ||
| : green_ctx_(other.green_ctx_), |
There was a problem hiding this comment.
Nit: can all done with member initializers using std::exchange
[ghstack-poisoned]
| // Convert to regular context | ||
| C10_CUDA_DRIVER_CHECK( | ||
| c10::cuda::DriverAPI::get()->cuCtxFromGreenCtx_(&context_, green_ctx_)); | ||
| TORCH_INTERNAL_ASSERT( |
There was a problem hiding this comment.
is it possible that C10_CUDA_DRIVER_CHECK is successful but conversion still fails? It should still be TORCH_CHECK in this case though, right? It's not "report a bug to pytorch", it's "something's wrong with your system"
| c10::cuda::DriverAPI::get()->cuCtxPushCurrent_(context_)); | ||
| } | ||
| // currently hardcoes the new green context to use the default stream | ||
| // TODO(eqy): consider creating a new stream if e.g., it allows interop |
There was a problem hiding this comment.
but it would still be as confusing to get 0 for torch.cuda.current_stream() which I think would happen now?
torch/csrc/cuda/green_context.h
Outdated
| cudaEvent_t ev; | ||
| C10_CUDA_CHECK(cudaEventCreate(&ev)); | ||
| C10_CUDA_CHECK(cudaEventRecord(ev, NULL)); |
There was a problem hiding this comment.
| cudaEvent_t ev; | |
| C10_CUDA_CHECK(cudaEventCreate(&ev)); | |
| C10_CUDA_CHECK(cudaEventRecord(ev, NULL)); | |
| CUDAEvent ev; | |
| ev.record(0); |
Did you mean to hardcode NULL stream here, or should it be current_stream?
Low-level PyTorch APIs should be usable/stable enough at this point but we might move the underlying driver API usage a bit from here...
Built on top of @drisspg 's branch
cc @ptrblck @msaroufim @jerryzh168