feat(inference): multi-GPU auto-layout and cross-GPU routing (PR2)

acestep/api/startup_llm_init.py

@@ -23,6 +23,7 @@ def initialize_llm_at_startup(
     get_model_name: Callable[[str], str],
     ensure_model_downloaded: Callable[[str, str], str],
     env_bool: Callable[[str, bool], bool],
+    dit_handler: Any = None,
 ) -> None:
     """Initialize LLM model according to GPU config and environment overrides."""
 
@@ -73,6 +74,9 @@ def initialize_llm_at_startup(
 
         lm_backend = resolve_lm_backend(os.getenv("ACESTEP_LM_BACKEND"), gpu_config)
         lm_device = os.getenv("ACESTEP_LM_DEVICE", device)
+        device_map = getattr(dit_handler, "device_map", None) if dit_handler is not None else None
+        if device_map is not None and device_map.lm is not None:
+            lm_device = device_map.lm
         lm_offload_env = os.getenv("ACESTEP_LM_OFFLOAD_TO_CPU")
         lm_offload = env_bool("ACESTEP_LM_OFFLOAD_TO_CPU", False) if lm_offload_env is not None else offload_to_cpu
 

acestep/api/startup_model_init.py

@@ -85,6 +85,7 @@ def do_model_initialization(
         compile_model=compile_model,
         offload_to_cpu=offload_to_cpu,
         offload_dit_to_cpu=offload_dit_to_cpu,
+        gpu_mapping=os.getenv("ACESTEP_GPU_MAPPING"),
     )
     if not ok:
         app.state._init_error = status_msg
@@ -157,6 +158,7 @@ def do_model_initialization(
         get_model_name=get_model_name,
         ensure_model_downloaded=ensure_model_downloaded,
         env_bool=env_bool,
+        dit_handler=handler,
     )
 
     print("[API Server] All models initialized successfully!")

acestep/core/generation/handler/audio_codes.py

@@ -56,7 +56,8 @@ def _decode_audio_codes_to_latents(self, code_str: str) -> Optional[torch.Tensor
         with self._load_model_context("model"):
             quantizer = self.model.tokenizer.quantizer
             detokenizer = self.model.detokenizer
-            indices = torch.tensor(code_ids, device=self.device, dtype=torch.long)
+            dit_device = self._get_component_device("model")
+            indices = torch.tensor(code_ids, device=dit_device, dtype=torch.long)
             indices = indices.unsqueeze(0).unsqueeze(-1)
 
             quantized = quantizer.get_output_from_indices(indices)
@@ -83,7 +84,11 @@ def convert_src_audio_to_codes(self, audio_file) -> str:
                         return "❌ Audio file appears to be silent"
                     latents = self._encode_audio_to_latents(processed_audio)
 
-                attention_mask = torch.ones(latents.shape[0], dtype=torch.bool, device=self.device)
+                attention_mask = torch.ones(
+                    latents.shape[0],
+                    dtype=torch.bool,
+                    device=self._get_component_device("model"),
+                )
                 with self._load_model_context("model"):
                     hidden_states = latents.unsqueeze(0)
                     _, indices, _ = self.model.tokenize(

acestep/core/generation/handler/conditioning_embed.py

@@ -56,7 +56,9 @@ def _ensure_latent_3d(z: torch.Tensor) -> torch.Tensor:
                         refer_audio = _normalize_audio_2d(refer_audio)
                         with torch.inference_mode():
                             refer_audio_latent = self.tiled_encode(refer_audio, offload_latent_to_cpu=True)
-                        refer_audio_latent = refer_audio_latent.to(self.device).to(self.dtype)
+                        refer_audio_latent = refer_audio_latent.to(
+                            self._get_component_device("dit")
+                        ).to(self.dtype)
                         if refer_audio_latent.dim() == 2:
                             refer_audio_latent = refer_audio_latent.unsqueeze(0)
                         refer_audio_latent = _ensure_latent_3d(refer_audio_latent.transpose(1, 2))
@@ -65,16 +67,28 @@ def _ensure_latent_3d(z: torch.Tensor) -> torch.Tensor:
                     refer_audio_order_mask.append(batch_idx)
 
         refer_audio_latents = torch.cat(refer_audio_latents, dim=0)
-        refer_audio_order_mask = torch.tensor(refer_audio_order_mask, device=self.device, dtype=torch.long)
+        dit_device = self._get_component_device("dit")
+        refer_audio_latents = refer_audio_latents.to(dit_device).to(self.dtype)
+        refer_audio_order_mask = torch.tensor(
+            refer_audio_order_mask, device=dit_device, dtype=torch.long
+        )
         return refer_audio_latents, refer_audio_order_mask
 
     def infer_text_embeddings(self, text_token_idss):
-        """Infer text-token embeddings via text encoder."""
+        """Infer text-token embeddings via text encoder.
+
+        Token ids are often built on the DiT device; move them onto the
+        text-encoder device so split maps like ``text_encoder:3`` work.
+        """
+        encoder_device = self._get_component_device("text_encoder")
+        text_token_idss = text_token_idss.to(encoder_device)
         with torch.inference_mode():
             return self.text_encoder(input_ids=text_token_idss, lyric_attention_mask=None).last_hidden_state
 
     def infer_lyric_embeddings(self, lyric_token_ids):
         """Infer lyric-token embeddings via text encoder embedding table."""
+        encoder_device = self._get_component_device("text_encoder")
+        lyric_token_ids = lyric_token_ids.to(encoder_device)
         with torch.inference_mode():
             return self.text_encoder.embed_tokens(lyric_token_ids)
 
@@ -124,6 +138,22 @@ def preprocess_batch(self, batch) -> Tuple:
 
         repaint_mask = batch.get("repaint_mask", None)
 
+        dit_device = self._get_component_device("dit")
+        text_hidden_states = text_hidden_states.to(dit_device)
+        lyric_hidden_states = lyric_hidden_states.to(dit_device)
+        text_attention_mask = text_attention_mask.to(dit_device)
+        lyric_attention_mask = lyric_attention_mask.to(dit_device)
+        src_latents = src_latents.to(dit_device)
+        target_latents = target_latents.to(dit_device)
+        chunk_mask = chunk_mask.to(dit_device)
+        is_covers = is_covers.to(dit_device) if isinstance(is_covers, torch.Tensor) else is_covers
+        if precomputed_lm_hints_25hz is not None:
+            precomputed_lm_hints_25hz = precomputed_lm_hints_25hz.to(dit_device)
+        if non_cover_text_hidden_states is not None:
+            non_cover_text_hidden_states = non_cover_text_hidden_states.to(dit_device)
+        if non_cover_text_attention_masks is not None:
+            non_cover_text_attention_masks = non_cover_text_attention_masks.to(dit_device)
+
         return (
             keys,
             text_inputs,

acestep/core/generation/handler/conditioning_embed_test.py

@@ -24,12 +24,21 @@ def embed_tokens(self, token_ids):
 class _Host(ConditioningEmbedMixin):
     """Minimal host implementing ConditioningEmbedMixin dependencies."""
 
-    def __init__(self):
-        self.device = "cpu"
+    def __init__(self, text_encoder_device="cpu", dit_device="cpu"):
+        self.device = dit_device
         self.dtype = torch.float32
         self.silence_latent = torch.zeros(1, 128, 6, dtype=torch.float32)
         self.text_encoder = _FakeTextEncoder()
         self.tiled_encode_calls = 0
+        self._devices = {
+            "dit": dit_device,
+            "vae": dit_device,
+            "text_encoder": text_encoder_device,
+            "lm": dit_device,
+        }
+
+    def _get_component_device(self, component):
+        return self._devices[component]
 
     def _ensure_silence_latent_on_device(self):
         return None
@@ -111,6 +120,25 @@ def test_preprocess_batch_returns_expected_tuple_shape(self):
         self.assertEqual(result[0], ["k1", "k2"])
         self.assertEqual(result[3].shape, (2, 128, 6))
 
+    def test_infer_text_embeddings_moves_ids_to_text_encoder_device(self):
+        """Token ids on DiT device must be moved onto a remote text encoder."""
+        if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
+            self.skipTest("needs 2+ CUDA devices")
+
+        seen = {}
+
+        class _DeviceCheckingEncoder(_FakeTextEncoder):
+            def __call__(self, input_ids, lyric_attention_mask=None):
+                seen["device"] = str(input_ids.device)
+                return super().__call__(input_ids, lyric_attention_mask)
+
+        host = _Host(text_encoder_device="cuda:1", dit_device="cuda:0")
+        host.text_encoder = _DeviceCheckingEncoder()
+        ids = torch.ones(1, 4, dtype=torch.long, device="cuda:0")
+        out = host.infer_text_embeddings(ids)
+        self.assertEqual(seen["device"], "cuda:1")
+        self.assertEqual(out.shape, (1, 4, 6))
+
 
 if __name__ == "__main__":
     unittest.main()

acestep/core/generation/handler/generate_music_decode.py

@@ -9,6 +9,7 @@
 from loguru import logger
 
 from acestep.gpu_config import get_effective_free_vram_gb
+from acestep.device_map import cuda_device_index, is_cuda_device
 
 
 class GenerateMusicDecodeMixin:
@@ -129,7 +130,12 @@ def _decode_generate_music_pred_latents(
         with torch.inference_mode():
             with self._load_model_context("vae"):
                 pred_latents_cpu = pred_latents.detach().cpu()
-                pred_latents_for_decode = pred_latents.transpose(1, 2).contiguous().to(self.vae.dtype)
+                vae_component_device = self._get_component_device("vae")
+                pred_latents_for_decode = (
+                    pred_latents.transpose(1, 2)
+                    .contiguous()
+                    .to(device=vae_component_device, dtype=self.vae.dtype)
+                )
                 del pred_latents
                 self._empty_cache()
 
@@ -140,7 +146,7 @@ def _decode_generate_music_pred_latents(
                 )
                 using_mlx_vae = self.use_mlx_vae and self.mlx_vae is not None
                 vae_cpu = False
-                vae_device = None
+                vae_restore_device = None
                 if not using_mlx_vae:
                     vae_cpu = os.environ.get("ACESTEP_VAE_ON_CPU", "0").lower() in ("1", "true", "yes")
                     if not vae_cpu:
@@ -150,7 +156,12 @@ def _decode_generate_music_pred_latents(
                                 "(unified memory), keeping VAE on MPS"
                             )
                         else:
-                            effective_free = get_effective_free_vram_gb()
+                            vae_cuda_index = (
+                                cuda_device_index(vae_component_device)
+                                if is_cuda_device(vae_component_device)
+                                else 0
+                            )
+                            effective_free = get_effective_free_vram_gb(vae_cuda_index)
                             logger.info(
                                 "[generate_music] Effective free VRAM before VAE decode: "
                                 f"{effective_free:.2f} GB"
@@ -163,7 +174,7 @@ def _decode_generate_music_pred_latents(
                                 vae_cpu = True
                     if vae_cpu:
                         logger.info("[generate_music] Moving VAE to CPU for decode (ACESTEP_VAE_ON_CPU=1)...")
-                        vae_device = next(self.vae.parameters()).device
+                        vae_restore_device = next(self.vae.parameters()).device
                         self.vae = self.vae.cpu()
                         pred_latents_for_decode = pred_latents_for_decode.cpu()
                         self._empty_cache()
@@ -186,9 +197,9 @@ def _decode_generate_music_pred_latents(
                         pred_wavs = decoder_output.sample
                         del decoder_output
                 finally:
-                    if vae_cpu and vae_device is not None:
+                    if vae_cpu and vae_restore_device is not None:
                         logger.info("[generate_music] Restoring VAE to original device after CPU decode path...")
-                        self.vae = self.vae.to(vae_device)
+                        self.vae = self.vae.to(vae_restore_device)
                     self._empty_cache()
                 logger.debug(
                     "[generate_music] After VAE decode: "

acestep/core/generation/handler/generate_music_decode_basic_test.py

@@ -0,0 +1,42 @@
+"""Tests for the default ``generate_music`` latent decode path."""
+
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from acestep.core.generation.handler.generate_music_decode_test_support import (
+    GENERATE_MUSIC_DECODE_MODULE,
+    DecodeTestHost,
+)
+
+
+class GenerateMusicDecodeBasicTests(unittest.TestCase):
+    """Verify successful decode timing and output behavior."""
+
+    def test_decode_pred_latents_updates_decode_time_and_returns_cpu_latents(self):
+        host = DecodeTestHost()
+        pred_latents = torch.ones(1, 4, 3)
+        time_costs = {"total_time_cost": 1.0}
+
+        def _progress(value, desc=None):
+            host.progress_calls.append((value, desc))
+
+        with patch.object(GENERATE_MUSIC_DECODE_MODULE.time, "time", side_effect=[10.0, 11.5]):
+            pred_wavs, pred_latents_cpu, updated_costs = host._decode_generate_music_pred_latents(
+                pred_latents=pred_latents,
+                progress=_progress,
+                use_tiled_decode=False,
+                time_costs=time_costs,
+            )
+
+        self.assertEqual(tuple(pred_wavs.shape), (1, 2, 8))
+        self.assertEqual(pred_latents_cpu.device.type, "cpu")
+        self.assertAlmostEqual(updated_costs["vae_decode_time_cost"], 1.5, places=6)
+        self.assertAlmostEqual(updated_costs["total_time_cost"], 2.5, places=6)
+        self.assertAlmostEqual(updated_costs["offload_time_cost"], 0.25, places=6)
+        self.assertEqual(host.progress_calls[0][0], 0.8)
+
+
+if __name__ == "__main__":
+    unittest.main()

acestep/core/generation/handler/generate_music_decode_cpu_offload_test.py

@@ -0,0 +1,100 @@
+"""Tests for CPU-offload and VAE restoration in ``generate_music`` decode."""
+
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from acestep.core.generation.handler.generate_music_decode_test_support import (
+    GENERATE_MUSIC_DECODE_MODULE,
+    DecodeTestHost,
+    FakeDecodeOutput,
+    FakeVae,
+)
+
+
+class GenerateMusicDecodeCpuOffloadTests(unittest.TestCase):
+    """Verify CPU-offload decode paths and device restoration."""
+
+    def test_decode_pred_latents_restores_vae_device_on_decode_error(self):
+        class FailingVae(FakeVae):
+            def __init__(self):
+                super().__init__()
+                self.cpu_calls = 0
+                self.to_calls = []
+
+            def decode(self, latents: torch.Tensor):
+                _ = latents
+                raise RuntimeError("decode failed")
+
+            def cpu(self):
+                self.cpu_calls += 1
+                return self
+
+            def to(self, *args, **kwargs):
+                self.to_calls.append((args, kwargs))
+                return self
+
+        class FailingHost(DecodeTestHost):
+            def __init__(self):
+                super().__init__()
+                self.use_mlx_vae = False
+                self.mlx_vae = None
+                self.vae = FailingVae()
+                self.empty_cache_calls = 0
+
+            def _empty_cache(self):
+                self.empty_cache_calls += 1
+
+        host = FailingHost()
+        with patch.dict(GENERATE_MUSIC_DECODE_MODULE.os.environ, {"ACESTEP_VAE_ON_CPU": "1"}, clear=False):
+            with self.assertRaisesRegex(RuntimeError, "decode failed"):
+                host._decode_generate_music_pred_latents(
+                    pred_latents=torch.ones(1, 4, 3),
+                    progress=None,
+                    use_tiled_decode=False,
+                    time_costs={"total_time_cost": 1.0},
+                )
+
+        self.assertEqual(host.vae.cpu_calls, 1)
+        self.assertEqual(len(host.vae.to_calls), 1)
+        self.assertGreaterEqual(host.empty_cache_calls, 2)
+
+    def test_decode_pred_latents_does_not_restore_latents_to_gpu_after_successful_cpu_decode(self):
+        class SuccessVae(FakeVae):
+            def __init__(self):
+                super().__init__()
+                self.vae_to_calls = []
+
+            def decode(self, latents: torch.Tensor):
+                return FakeDecodeOutput(torch.ones(latents.shape[0], 2, 8))
+
+            def cpu(self):
+                return self
+
+            def to(self, *args, **kwargs):
+                self.vae_to_calls.append(args[0] if args else kwargs)
+                return self
+
+        class SuccessHost(DecodeTestHost):
+            def __init__(self):
+                super().__init__()
+                self.use_mlx_vae = False
+                self.mlx_vae = None
+                self.vae = SuccessVae()
+
+        host = SuccessHost()
+        with patch.dict(GENERATE_MUSIC_DECODE_MODULE.os.environ, {"ACESTEP_VAE_ON_CPU": "1"}, clear=False):
+            pred_wavs, _cpu_latents, _costs = host._decode_generate_music_pred_latents(
+                pred_latents=torch.ones(1, 4, 3),
+                progress=None,
+                use_tiled_decode=False,
+                time_costs={"total_time_cost": 1.0},
+            )
+
+        self.assertEqual(len(host.vae.vae_to_calls), 1)
+        self.assertEqual(tuple(pred_wavs.shape), (1, 2, 8))
+
+
+if __name__ == "__main__":
+    unittest.main()