Skip to content

compat: expose torch.cuda.streams for torch 2.11 MUSA - #98

Merged
yeahdongcn merged 2 commits into
mainfrom
codex/torch-musa-streams-compat
Aug 7, 2026
Merged

yeahdongcn merged 2 commits into
mainfrom
codex/torch-musa-streams-compat

Conversation

@yeahdongcn

Copy link
Copy Markdown
Collaborator

Summary

Expose the MUSA stream module through the CUDA-compatible torch.cuda.streams path.

PyTorch 2.11 Dynamo guards access torch.cuda.streams.Stream. torchada 0.1.78 redirects this path to torch_musa, but torch_musa 2.11 exposes the implementation at torch_musa.core.stream and does not export torch_musa.streams, causing compiled vLLM startup to fail with:

AttributeError: module 'torch_musa' has no attribute 'streams'

The shim maps both the attribute and sys.modules["torch.cuda.streams"] to torch_musa.core.stream.

Validation

  • torch.cuda.streams is torch_musa.core.stream: PASS on MUSA 3.3.5 / torch 2.11.0.post1+musa5.2.0.
  • Candidate vLLM compiled startup: PASS after the shim; the previous Dynamo guard failure disappeared.
  • vLLM dense A/B on the same 4-GPU host: candidate performance remained positive across BS 1/4/16/64.

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
@yeahdongcn
yeahdongcn marked this pull request as ready for review August 7, 2026 06:58
@yeahdongcn
yeahdongcn merged commit 71bfc5d into main Aug 7, 2026
@yeahdongcn
yeahdongcn deleted the codex/torch-musa-streams-compat branch August 13, 2026 12:16
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.

1 participant