Skip to content

fix(musa): expose device index helper for TorchDynamo - #108

Merged
yeahdongcn merged 5 commits into
MooreThreads:mainfrom
wukecong-mt:fix/torchdynamo-device-index
Aug 28, 2026
Merged

yeahdongcn merged 5 commits into
MooreThreads:mainfrom
wukecong-mt:fix/torchdynamo-device-index

Conversation

@wukecong-mt

Copy link
Copy Markdown
Contributor

Summary

Map torch.cuda._get_device_index to
torch_musa.core._utils._get_musa_device_index in the CUDA module wrapper.

Motivation

TorchDynamo models torch.cuda.device(...) by calling
torch.cuda._get_device_index(...).

Without the mapping, TorchAda redirects the lookup to the top-level
torch_musa module, where the helper does not exist:

InternalTorchDynamoError:
AttributeError: module 'torch_musa' has no attribute '_get_device_index' 

Reproduction

import torch
import torchada

x = torch.ones(4, device="musa")

@torch.compile(backend="eager")
def f(t):
    with torch.cuda.device(t.device.index):
        return t + 1

print(f(x))

@yeahdongcn yeahdongcn left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Please add a commit to bump version (replace all, should be 6 files).

Comment thread src/torchada/_patch.py Outdated
@yeahdongcn
yeahdongcn merged commit 90d8d2a into MooreThreads:main Aug 28, 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