Add graph-safe ops and Inductor heuristics

[froststeam]

Summary

Add the graph-safe MUSA compatibility required by CUDA-facing inference paths and bump torchada from 0.1.83 to 0.1.84:

• add MUSA PrivateUse1 overrides for aten::multinomial, aten::log, and aten::log_ on affected torch_musa releases
• reuse CUDA Inductor Triton template heuristics for MUSA matmul templates
• mirror MUSA_VISIBLE_DEVICES to CUDA_VISIBLE_DEVICES when both are present
• add focused correctness and CUDA-graph replay coverage for the new operators

torch_musa version boundary

The TorchAda graph-safe compatibility is temporary and is limited to torch_musa versions before 2.11.0.post2:

torch_musa < 2.11.0.post2  -> keep TorchAda compatibility patches
torch_musa >= 2.11.0.post2 -> use torch_musa/ATen native implementations

The existing _TORCH_MUSA_ACCELERATOR_FIX_VERSION = "2.11.0.post2" is reused for this decision. The version is read from the existing torch.musa.__version__ surface; no extra package-version helper or unconditional import torch_musa is added just for version detection.

On torch_musa >= 2.11.0.post2:

• load_cpp_ops() returns before building/loading TorchAda's C++ override extension
• the Inductor heuristic compatibility patch is skipped
• torch_musa/ATen owns log, log_, and multinomial dispatch directly
• no compile-time macro or version check is kept in musa_ops.mu

The currently published vLLM-MUSA image is still based on torch_musa 2.11.0.post1+musa5.2.0, so it remains on the compatibility side of this boundary.

Scope

This PR does not change global CUDA platform semantics:

• does not patch torch.cuda.is_available()
• does not fake torch.cuda.get_device_capability()
• does not globally rewrite torch.load map locations
• does not add a flash-attention import shim
• does not add SGLang-Omni-specific model logic

Testing

Focused validation was run with the PR source overlaid in the MUSA-enabled vLLM-MUSA image:

PYTHONPATH=src python -m pytest -q tests/test_cuda_patching.py \
  -k "VisibleDevicesEnv or InductorTemplateHeuristics or TestAcceleratorModuleWrapper or test_cpp_ops_loaded_on_musa"

36 passed, 202 deselected

Operator correctness and graph replay tests on the same post1 MUSA runtime:

PYTHONPATH=src python -m pytest -q tests/test_log.py tests/test_multinomial.py

8 passed

The version-boundary tests cover the existing post2 accelerator gate, post2 native Inductor registration behavior, local-version suffixes such as +musa5.2.0, and the post2+ no-C++-override behavior. git diff --check and Python syntax compilation also pass.

Downstream image validation

Downstream image	torch / torch_musa	Triton keys	Missing MUSA keys	Result
registry.mthreads.com/mcconline/inference/sglang:v0.5.12.post1-ph1-4.3.5-torch2.9.0-20260819	2.9.0 / 40305	11	0	SGLang import and MUSA matmul passed after installing the image's optional runtime dependencies
registry.mthreads.com/mcconline/inference/vllm:v0.24.0-ph1-5.2.0-torch2.11.0.post1-20260827	2.11.0.post1+musa5.2.0 / 50200	15	0	vLLM import, MUSA platform plugin loading, and MUSA matmul passed

A direct runtime probe in the latest image confirmed:

torch_musa = 2.11.0.post1+musa5.2.0
torch.musa.is_available() = True
torch.musa.device_count() = 1
aten::log PrivateUse1 kernel = present
aten::log_ PrivateUse1 kernel = present
aten::multinomial PrivateUse1 kernel = present

Request-level vLLM compile smoke

A request-level Qwen3-0.6B vLLM compile and GSM8K follow-up was run on MUSA with CompilationMode.VLLM_COMPILE (mode=3) using the compatible image registry.mthreads.com/mcconline/inference/vllm:v0.24.0-ph1-5.2.0-torch2.9.1.post1-20260805.

The service completed compilation, graph capture, a direct completion request, and the official 200-question GSM8K evaluator:

Accuracy: 0.380
Invalid responses: 0.000
All 200 requests returned HTTP 200

The server log contained no runtime exception, backend compiler failure, graph break, or eager fallback during evaluation.

Limitations

• No public torch_musa 2.11.0.post2 wheel/tag is available yet; post2 behavior is covered by version-boundary tests and the code path is fail-safe for unknown versions.
• The 20260819 and 20260824 Qwen3 startup attempts hit a pre-existing vLLM-MUSA is_neox signature mismatch, so request-level evaluation used the compatible 20260805 image.
• The request-level run validates the Python Inductor path and downstream service behavior; it should not be interpreted as C++ override coverage for a post2 wheel.

[yeahdongcn]

This duplicates a full operator implementation at the MUSA kernel layer. ATen schemas/signatures, generator/RNG semantics, dtype handling, and edge-case behavior can change across torch versions, so this can keep compiling while silently diverging from upstream. Could we first reuse the torch_musa/ATen implementation and keep only a thin compatibility shim, or introduce an explicit versioned adapter? If a custom override is required, please add differential tests across the target torch versions (generator, replacement, invalid/NaN/negative inputs, empty and boundary shapes).

[yeahdongcn]

This copies the current CUDA heuristic registry only once during import. That will miss lazy/future registrations, and it relies on private registry/cache names and key shape. Also, copying a CUDA heuristic class under a musa key does not establish that its lowering/template/autotune path is MUSA-compatible. Could we move this compatibility to the registration/lookup boundary (or use an explicit versioned allowlist) and add a real torch.compile + Inductor/Triton MUSA smoke after lazy imports? Unsupported templates/torch versions should fail closed or fall back rather than silently appearing supported.

[lijing-mt]

Current registration process is:

import torch._inductor.template_heuristics.registry as registry
----register all torch native "triton::" heuristics from `torch/_inductor/template_heuristics/triton.py`

......

from torch._inductor.codegen.common import init_backend_registration
init_backend_registration()
----lazy register all musa heuristics from `torch_musa/_inductor/template_heuristics.py`

When we need to fully import custom heuristics, it is recommended to modify heuristic_registry after calling init_backend_registration().