Skip to content

Add CUDA memory pool allocation function aliases for MUSA compatibility - #61

Merged
yeahdongcn merged 4 commits into
MooreThreads:mainfrom
popsiclexu:xzx/mem_pool
May 7, 2026
Merged

yeahdongcn merged 4 commits into
MooreThreads:mainfrom
popsiclexu:xzx/mem_pool

Conversation

@popsiclexu

Copy link
Copy Markdown
Contributor
  • Add CUDA memory pool allocation function aliases (_cuda_beginAllocateCurrentThreadToPool, _cuda_endAllocateToPool, _cuda_releasePool) in the C++ extension, backed by MUSA allocator APIs, to support from torch.cuda.memory import ... usage on MUSA devices
  • Implement thread-based filtering in _musa_beginAllocateCurrentThreadToPool to match CUDA semantics (instead of the stream-based fallback from _patch.py)
  • Inject these functions into torch.musa.memory at C++ extension load time, so downstream code using CUDA memory pool APIs works transparently on MUSA

@augmentcode

augmentcode Bot commented May 6, 2026

Copy link
Copy Markdown
🤖 Augment PR Summary

Summary: This PR improves CUDA-API compatibility on MUSA by exposing CUDA-named memory pool allocation helpers through torchada’s C++ extension.

Changes:

  • Adds C++ bindings for _cuda_beginAllocateCurrentThreadToPool, _cuda_endAllocateToPool, and _cuda_releasePool backed by MUSA caching allocator APIs.
  • Implements thread-based filtering for “current thread” semantics when allocating to a pool.
  • Injects the new functions into torch.musa.memory when the extension loads so torch.cuda.memory imports work under torchada’s CUDA->MUSA redirection.

Technical Notes: Injection happens after the extension is successfully built/loaded and marked as loaded, and is intended to make downstream CUDA memory pool API calls operate transparently on MUSA.

🤖 Was this summary useful? React with 👍 or 👎

@augmentcode augmentcode Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review completed. 2 suggestions posted.

Fix All in Augment

Comment augment review to trigger a new review at any time.

Comment thread src/torchada/_cpp_ops.py Outdated
Comment thread src/torchada/_cpp_ops.py Outdated
@popsiclexu
popsiclexu force-pushed the xzx/mem_pool branch 2 times, most recently from 0fb07bb to 28d1545 Compare May 6, 2026 11:50
Comment thread src/torchada/_cpp_ops.py
Comment thread src/torchada/_patch.py Outdated
Comment thread src/torchada/_patch.py Outdated
yeahdongcn added 3 commits May 7, 2026 14:52
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
@yeahdongcn
yeahdongcn merged commit e4054d9 into MooreThreads:main May 7, 2026
yeahdongcn added a commit that referenced this pull request Oct 10, 2026
The English and Chinese READMEs have not kept up with the May-October work.
This documents what merged, moves the version-gated shims into one table, and
refreshes the measured numbers that had gone stale.

Feature table
- CUDA memory-pool APIs, `torch.cuda.streams`, CUDA-graph executable rotation,
  `torch.cuda._get_device_index`, `get_memory_info()`, and the FlashAttention
  provider shims (#61, #98, #103, #106, #108, #115)
- the "What Works" table goes back to one line per feature; the paragraph-sized
  `log_` / `isfinite` / `out_dtype` cells move into the new section below

New "torch_musa Compatibility" section
- one table of every version-gated shim with the release it is installed on:
  the four `< 2.11.0.post2` patches (#106, #113, #124), the `< 2.13.0`
  `mm`/`bmm` `out_dtype=` backport (#116), the stable-ABI header backport (#86,
  #96), asynchronous `isfinite` (#120), and `torch.cuda.streams` (#98)

New "Environment Variables" section
- the graph-rotation knobs (#72), `TORCHADA_PLATFORM`, the C++ operator-override
  switches (#61, #128), and the two variables that were already documented

Corrected and extended details
- torch.compile: FX `device` builtin (#124), Dynamo's device-index helper (#108),
  `MUSA_VISIBLE_DEVICES` mirroring (#106)
- C++ extensions: nested `<torch/cuda.h>` porting (#95), stable-ABI
  `STABLE_TORCH_LIBRARY_IMPL` rekeying and stream helpers (#100), torch 2.6+
  `include_paths`/`library_paths` signatures (#121), stale JIT build locks (#128)
- MoE tables are generated from checked-in recipes (#115)
- unsupported CUDA runtime APIs as no-ops (#65)
- Performance: replace the 0.1.94 / torch_musa 2.7.1 numbers with the checked-in
  0.1.95 / 2.11.0.post2 entry, and stop claiming every fast path is under 200ns
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants