[PR #27] fix: implement ELF GOT patching for CUDA hooks on Linux/WSL2 #70

Open
opened 2026-10-26 22:24:18 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-aimdo/pull/27
Author: @WASasquatch
Created: 3/20/2026
Status: 🔄 Open

Base: master ← Head: fix/linux-wsl2-cuda-hooks


📝 Commits (1)

  • e7f2fd2 fix: implement ELF GOT patching for CUDA hooks on Linux/WSL2

📊 Changes

5 files changed (+390 additions, -49 deletions)

View changed files

📝 .github/workflows/build-wheels.yml (+1 -1)
📝 build-linux-docker (+1 -1)
➕ src-posix/cuda-hook.c (+379 -0)
📝 src/plat.h (+3 -2)
📝 src/pyt-cu-plug-alloc-async.c (+6 -45)

📄 Description

On Linux, libcudart resolves driver API symbols via dlsym at runtime, so the no-op stubs left GOT entries unpatched and VRAM tracking broken.

This commit:

  • Adds src-posix/cuda-hook.c with ELF GOT/PLT patching that hooks both CUDA runtime API (cudaMalloc/cudaFree/cudaMallocAsync/cudaFreeAsync) and driver API (cuMemAlloc_v2/cuMemFree_v2/cuMemAllocAsync/cuMemFreeAsync)
  • Runtime hooks call through real runtime API preserving CUDA memory pools
  • Self-contained allocation tracking via thread-safe hash table
  • Resolves runtime symbols via RTLD_DEFAULT (handles versioned sonames)
  • Leaves GOT pages R+W after patching for lazy resolution compatibility
  • Adds -ldl to build scripts for dlopen/dlsym/dl_iterate_phdr
  • Removes dead symbol-interposition overrides from pyt-cu-plug-alloc-async.c

Tested on RTX 4090 under WSL2 (Ubuntu 24.04) with PyTorch 2.10+cu128.

Contribution Agreement

  • I agree that my contributions are licensed under the GPLv3.
  • I grant Comfy Org the rights to relicense these contributions as outlined in CONTRIBUTING.md.

🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/Comfy-Org/comfy-aimdo/pull/27 **Author:** [@WASasquatch](https://github.com/WASasquatch) **Created:** 3/20/2026 **Status:** 🔄 Open **Base:** `master` ← **Head:** `fix/linux-wsl2-cuda-hooks` --- ### 📝 Commits (1) - [`e7f2fd2`](https://github.com/Comfy-Org/comfy-aimdo/commit/e7f2fd2d9db8aaac31d3f0b5398b3149449cddd2) fix: implement ELF GOT patching for CUDA hooks on Linux/WSL2 ### 📊 Changes **5 files changed** (+390 additions, -49 deletions) <details> <summary>View changed files</summary> 📝 `.github/workflows/build-wheels.yml` (+1 -1) 📝 `build-linux-docker` (+1 -1) ➕ `src-posix/cuda-hook.c` (+379 -0) 📝 `src/plat.h` (+3 -2) 📝 `src/pyt-cu-plug-alloc-async.c` (+6 -45) </details> ### 📄 Description On Linux, libcudart resolves driver API symbols via dlsym at runtime, so the no-op stubs left GOT entries unpatched and VRAM tracking broken. This commit: - Adds src-posix/cuda-hook.c with ELF GOT/PLT patching that hooks both CUDA runtime API (cudaMalloc/cudaFree/cudaMallocAsync/cudaFreeAsync) and driver API (cuMemAlloc_v2/cuMemFree_v2/cuMemAllocAsync/cuMemFreeAsync) - Runtime hooks call through real runtime API preserving CUDA memory pools - Self-contained allocation tracking via thread-safe hash table - Resolves runtime symbols via RTLD_DEFAULT (handles versioned sonames) - Leaves GOT pages R+W after patching for lazy resolution compatibility - Adds -ldl to build scripts for dlopen/dlsym/dl_iterate_phdr - Removes dead symbol-interposition overrides from pyt-cu-plug-alloc-async.c Tested on RTX 4090 under WSL2 (Ubuntu 24.04) with PyTorch 2.10+cu128. ### Contribution Agreement - [ ] I agree that my contributions are licensed under the GPLv3. - [ ] I grant **Comfy Org** the rights to relicense these contributions as outlined in [CONTRIBUTING.md](./CONTRIBUTING.md). --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
Sign in to join this conversation.
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
Comfy-Org/comfy-aimdo#70
No description provided.