Add graph-safe ops and Inductor heuristics

[froststeam]

Summary

Add the minimal torchada compatibility needed by CUDA-compatible MUSA model paths and bump the package version from 0.1.83 to 0.1.84:

• register MUSA PrivateUse1 overrides for graph-safe aten::multinomial, aten::log, and aten::log_
• reuse CUDA Inductor template heuristics for MUSA Triton matmul templates, including triton::mm and triton::mm with addmm
• update the documented/installable torchada version to 0.1.84
• add focused MUSA tests for the new operator overrides, including CUDA graph replay coverage

Scope

This PR keeps torchada behavior narrow and avoids changing global platform semantics:

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

The intent is to preserve existing CUDA-facing behavior while adding the missing MUSA-safe operator and Inductor coverage required by downstream graph-enabled paths.

Testing

uvx pre-commit run --files \
  pyproject.toml \
  src/torchada/__init__.py \
  tests/test_platform.py \
  README.md \
  README_CN.md \
  src/torchada/_patch.py \
  src/torchada/csrc/musa_ops.mu \
  tests/test_log.py \
  tests/test_multinomial.py

Result: passed.

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

Result on a MUSA-enabled runtime:

18 passed

Runtime probe on the same MUSA-enabled runtime:

torch.cuda.is_available() == False
torch.musa.is_available() == True
("triton::mm", "musa", None) registered in Inductor heuristic registry
("triton::mm", "musa", "addmm") registered in Inductor heuristic registry
aten::multinomial PrivateUse1 override loaded from torchada MUSA ops

Additional SGLang-Omni integration smoke with the latest SGLang, SGLang-Omni, and torchada changes:

Model	Request	Result
Qwen/Qwen3-ASR-1.7B	/v1/audio/transcriptions with a WAV input	Passed; returned text transcription
Qwen/Qwen3-TTS-12Hz-0.6B-Base	/v1/audio/speech with graph enabled and reduced capture range	Passed; generated WAV audio
Qwen/Qwen3-Omni-30B-A3B-FP8	image chat completion with graph enabled	Passed; returned the expected image answer
MiniMaxAI/MiniMax-Music3	/v1/audio/speech music generation	Passed; generated stereo WAV audio

MUSA Triton heuristic regression coverage

The review follow-up adds focused unit coverage for _patch_inductor_template_heuristics using representative keys observed in downstream inference images (triton::bmm, triton::mm, and triton::mm with addmm).

Covered behavior:

• copies CUDA Triton heuristic entries to the corresponding MUSA key
• preserves an explicitly registered MUSA heuristic instead of overwriting it
• ignores non-Triton and non-CUDA entries
• safely ignores malformed and non-tuple registry keys
• clears _HEURISTIC_CACHE only when the registry changes
• remains idempotent on repeated application
• performs no registry or cache mutation on non-MUSA platforms
• makes MUSA_VISIBLE_DEVICES authoritative when both MUSA and CUDA visibility variables are set

Focused unit test:

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

Result:

7 passed, 230 deselected

Image-level regression matrix with the PR source overlaid:

Downstream image	PyTorch / MUSA	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 passed after installing the image's optional runtime dependencies; MUSA matmul passed
registry.mthreads.com/mcconline/inference/vllm:v0.24.0-ph1-5.2.0-torch2.11.0.post1-20260824	2.11.0.post1+musa5.2.0 / 50200	15	0	vLLM import and MUSA platform plugin loading passed; MUSA matmul passed

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

Applied the MUSA Qwen3 QK-RoPE-KV pre-split fusion to 28 layers.
Using cache directory: .../torch_compile_cache/.../backbone for vLLM's torch.compile
Dynamo bytecode transform time: 5.33 s
Compiling a graph for compile range (1, 8192) takes 3.23 s
torch.compile and initial profiling/warmup run together took 53.06 s in total
Graph capturing finished in 638 secs, took 6.16 GiB

A direct /v1/completions request returned readable text and the correct numerical answer. The official vLLM GSM8K endpoint evaluator then completed 200 questions with 5-shot prompting, temperature 0, and 256 maximum output tokens:

Accuracy: 0.380
Invalid responses: 0.000
Total latency: 77.376 s
Questions per second: 2.585
Total output tokens: 20288
Output tokens per second: 262.200

All 200 requests returned HTTP 200; the server log contained no runtime exception, backend compiler failure, graph break, or eager fallback during evaluation.

[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.