[Bugfix][MiMo-Audio] Align code2wav decode with tokenizer device

[smartDream-chao]

What broke

When MIMO_AUDIO_TOKENIZER_DEVICE=cpu (or the tokenizer worker otherwise lives on CPU), Stage-1 code2wav crashes during VQ decode:

RuntimeError: Expected all tensors to be on the same device, but got index is on cuda:0, different from other tensors on cpu
  (when checking argument in method wrapper_CUDA__index_select)

Additionally, building torchaudio.transforms.MelSpectrogram while the default CUDA device is active can mix CPU/CUDA tensors on newer PyTorch.

Repro steps

1. Single-GPU RTX 5090D (Blackwell) setup with MiMo-Audio offline TTS.
2. Export MIMO_AUDIO_TOKENIZER_DEVICE=cpu and MIMO_AUDIO_TOKENIZER_CUDA_GRAPH=0.
3. Run examples/offline_inference/mimo_audio/end2end.py with tts_sft.
4. Observe Stage-1 fatal error in decode_vq → F.embedding.

Root cause

• Batch / streaming code2wav paths moved audio_codes (and input_lengths) to the stage device (self.device, CUDA) while tokenizer weights stayed on the tokenizer device (CPU).
• Stage-1 did not honor MIMO_AUDIO_TOKENIZER_DEVICE when constructing the tokenizer worker.
• MelSpectrogram init under a CUDA default device could fail or mix devices.

Fix

• Build MelSpectrogram under torch.device("cpu"), then move to the tokenizer device when needed.
• Honor MIMO_AUDIO_TOKENIZER_DEVICE when creating the Stage-1 tokenizer worker.
• Use tok_device = torch.device(self._tokenizer_service.device) for decode_vq / CUDA-graph decode / input_lengths.

Test plan

• Offline TTS (tts_sft) on 1x RTX 5090D 32GB with tokenizer on CPU; wav written under output_audio/tts_sft/.
• Default CUDA tokenizer path still loads (device env unset).
• CI / lint on this PR.

Notes

• Validated primarily with CPU tokenizer to keep Stage-1 VRAM headroom on a shared 5090D with Stage-0.
• Recipe / hardware knobs for Blackwell (e.g. TRITON_ATTN) will be a follow-up docs PR.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/model_integration.md.

Module owners: @tzhouam @gcanlin

@smartDream-chao, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[smartDream-chao]

Self-review

• Verified the Stage-1 crash: with MIMO_AUDIO_TOKENIZER_DEVICE=cpu, decode_vq / F.embedding failed because codes were moved to CUDA (self.device) while tokenizer weights stayed on CPU.
• Reviewed the three changes in mimo_audio_code2wav.py:
  1. Build MelSpectrogram under torch.device("cpu"), then move to the tokenizer device when needed.
  2. Honor MIMO_AUDIO_TOKENIZER_DEVICE when constructing the Stage-1 tokenizer worker.
  3. Batch / streaming code2wav decode uses tok_device from the tokenizer service for codes, CUDA-graph decode, and input_lengths.
• Validated offline tts_sft on 1x RTX 5090D 32GB (tokenizer on CPU); wav written under output_audio/tts_sft/.
• Confirmed the commit includes Signed-off-by for DCO.
• Scope is limited to one file; Blackwell recipe / deploy knobs left for a follow-up docs PR.

[smartDream-chao]

Friendly bump — this is still relevant on current main
(decode_vq still uses audio_codes.to(self.device) while the tokenizer
can live on CPU via MIMO_AUDIO_TOKENIZER_DEVICE).

Just rebased onto latest main so CI is fresh. Scope remains one
file (mimo_audio_code2wav.py); self-review + 5090D repro notes are above.

Would appreciate a look from the requested reviewers when you have a chance.
Thanks!

[linyueqian]

Request changes. The device placement fix is correct, but the new service attribute breaks seven of the PR's eight batch-decode tests before decoding. Indexed CUDA targets are also reduced to the current device.

Validation: static review at bc786fae of the full diff, batch-decode fixture, worker cache key, and CPU graph-capture guard. I did not run code from the untrusted fork.

[linyueqian]

The new self._tokenizer_service.device read breaks the batch-decode tests. Their fixture creates the service without a device attribute, so seven of eight cases raise AttributeError before decoding. Add device="cpu" to the fixture; the second decode method has the same dependency.

[linyueqian]

An indexed target such as cuda:1 is reduced to cuda here, so the worker uses the process's current CUDA device instead of device 1. Preserve the index through the cache key and constructor, or reject indexed values explicitly.

[smartDream-chao]

Thanks for the careful review @linyueqian — both points were real. Pushed 22990118.

1. Batch-decode fixture (AttributeError)

Confirmed: _minimal_model built the tokenizer service without a device
attribute, and tok_device is read before decoding, so only the empty-input
case survived. Added device="cpu" to the fixture; both decode methods share
it, so the streaming path is covered too.

2. Indexed CUDA targets

Went with preserving the index rather than rejecting it. _normalize_tokenizer_worker_cache_key
now keeps the full device string, and get_tokenizer_worker passes it to the
worker constructor, so cuda:1 no longer falls back to the process's current
device or shares an entry with cuda:0.

One extra detail: keeping the raw string would have let a bare cuda and
cuda:0 occupy two entries and load the tokenizer twice, which defeats the
function's stated purpose. So a bare cuda is resolved to
torch.cuda.current_device() when CUDA is available, and both spellings
collapse to one entry. Also widened the device annotations to
torch.device | str, which is what the code already accepted.

Added tests/model_executor/models/mimo_audio/test_mimo_audio_tokenizer_worker_cache.py
covering index retention (device and string spellings), cpu, bare-cuda
resolution, constructor propagation, distinct workers per index, and cache
reuse. Worker class is stubbed, so it is CPU-only and weight-free.

Happy to switch to explicit rejection of indexed values instead if you would
rather keep this PR's surface smaller.

[smartDream-chao]

Friendly re-review ping @linyueqian — both of your review points should be addressed now:

1. batch-decode fixture: added device="cpu" on the tokenizer service stub
2. indexed CUDA targets: cache key + worker constructor keep cuda:N; bare cuda resolves via torch.accelerator.current_device_index()

GitHub checks and Buildkite are green after the ready label. Would you mind taking another look / resolving the threads when you have a chance? Thanks!

[linyueqian]

Looks good. The fixture now supplies the tokenizer-service device, and indexed CUDA targets stay intact through both the cache key and worker construction. Both prior findings are resolved, and I found no new material issue.

Validation: exact-head static review at 581375d2 of the full patch, both follow-up commits, both decode paths, cache normalization, worker construction, focused regression tests, repository accelerator conventions, and git diff --check. All four per-PR Buildkite lanes are green at this head, which is the runtime evidence that matters most here since both findings concerned device selection and those lanes exercise it on general, AMD, Intel, and NPU. I did not execute the branch myself, so the approval rests on static review plus that lane coverage rather than on a local run.