[Bugfix][Hardware][Ascend] Fix Qwen3-TTS code predictor dtype and nested graph replay (#6639)

Signed-off-by: tanyadong <2209931449@qq.com>
Signed-off-by: tanyadong <yadongtan666@gmail.com>
Co-authored-by: tanyadong <2209931449@qq.com>

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -572,6 +572,53 @@ class TestCodePredictorPerRowGenerators:
 class TestCodePredictorModelDtype:
     """Test the inner model forward with different dtypes."""
 
+    @pytest.mark.parametrize(
+        ("device_type", "expect_fp32"),
+        [("cpu", False), ("npu", False), ("xpu", True), ("musa", True), ("cuda", True)],
+    )
+    def test_fp32_fallback_excludes_cpu_and_npu(
+        self,
+        mocker: MockerFixture,
+        loaded_target_classes,
+        device_type: str,
+        expect_fp32: bool,
+    ) -> None:
+        """The stability fallback preserves accelerator behavior except on NPU."""
+        _, _, _, code_predictor_model, _ = loaded_target_classes
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+
+        inputs = mocker.MagicMock()
+        inputs.dtype = torch.float16
+        inputs.device.type = device_type
+        fp32_inputs = mocker.MagicMock()
+        fp32_inputs.device.type = device_type
+        inputs.float.return_value = fp32_inputs
+        model = object.__new__(code_predictor_model)
+        torch.nn.Module.__init__(model)
+        model.rotary_emb = mocker.Mock(return_value=None)
+        model.layers = []
+        model.norm = mocker.Mock(side_effect=lambda hidden_states: hidden_states)
+        autocast = mocker.patch.object(
+            common_mod.torch.amp,
+            "autocast",
+            return_value=mocker.MagicMock(
+                __enter__=mocker.Mock(),
+                __exit__=mocker.Mock(return_value=False),
+            ),
+        )
+
+        code_predictor_model.forward(model, inputs, mocker.MagicMock())
+
+        if expect_fp32:
+            inputs.float.assert_called_once_with()
+        else:
+            inputs.float.assert_not_called()
+        autocast.assert_called_once_with(
+            device_type,
+            enabled=expect_fp32,
+            dtype=torch.float32,
+        )
+
     def test_model_forward_float16(self, loaded_target_classes) -> None:
         """Inner model forward should work in float16."""
         _, _, _, code_predictor_model, _ = loaded_target_classes
@@ -601,6 +648,67 @@ class TestCodePredictorModelDtype:
         assert output.shape == (bsz, seq_len, 32)
 
 
+class TestCodePredictorGraphReplay:
+    """Test nested device graph handling for the shared code predictor."""
+
+    @pytest.mark.parametrize("is_capturing", [False, True])
+    def test_npu_outer_capture_skips_inner_graph_replay(
+        self,
+        mocker: MockerFixture,
+        loaded_target_classes,
+        is_capturing: bool,
+    ) -> None:
+        _, _, code_predictor_wrapper, _, code_predictor_wrapper_config = loaded_target_classes
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+
+        predictor = object.__new__(code_predictor_wrapper)
+        torch.nn.Module.__init__(predictor)
+        predictor._num_groups = 3
+        predictor._model_dtype = torch.float32
+        predictor._setup_compile = mocker.Mock()
+        predictor._padded_bsz = mocker.Mock(side_effect=lambda bsz: bsz)
+        predictor._ensure_buffers = mocker.Mock()
+        predictor._proj_buf = torch.zeros(1, 4, 4)
+        predictor.small_to_mtp_projection = torch.nn.Identity()
+        regular_output = torch.zeros(1, 4, 4)
+        predictor._compiled_model_fwd = mocker.Mock(return_value=regular_output)
+        predictor._lm_heads_list = [
+            torch.nn.Linear(4, 8, bias=False),
+            torch.nn.Linear(4, 8, bias=False),
+        ]
+        predictor._codec_embeds_list = [torch.nn.Embedding(8, 4)]
+        predictor._wrapper_config = code_predictor_wrapper_config(
+            sampling_mode="per_call",
+            return_proj_buf=False,
+        )
+        predictor._prefix_graphs_enabled = False
+        predictor._bucket_pos_ids = {1: torch.arange(4).unsqueeze(0)}
+        graph = mocker.Mock()
+        graph_output = torch.zeros(1, 4, 4)
+        predictor._device_graphs = {1: (graph, graph_output)}
+
+        mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=True)
+        npu_mock = mocker.MagicMock()
+        npu_mock.is_current_stream_capturing.return_value = is_capturing
+        mocker.patch.object(common_mod.torch, "npu", npu_mock, create=True)
+
+        result = predictor(
+            layer0_code=torch.zeros(1, dtype=torch.long),
+            layer0_embed=torch.zeros(1, 4),
+            last_talker_hidden=torch.zeros(1, 4),
+            do_sample=False,
+        )
+
+        assert result.shape == (1, 3)
+        npu_mock.is_current_stream_capturing.assert_called_once_with()
+        if is_capturing:
+            graph.replay.assert_not_called()
+            assert predictor._compiled_model_fwd.call_count == 2
+        else:
+            assert graph.replay.call_count == 2
+            predictor._compiled_model_fwd.assert_not_called()
+
+
 class TestCodePredictorWrapperConfig:
     """Test wrapper configuration for different models."""
 

vllm_omni/model_executor/models/common/qwen3_code_predictor.py

@@ -439,10 +439,11 @@ class CodePredictorBaseModel(nn.Module):
         # native bf16 support (Turing, Volta).  The RMSNorm and RoPE
         # layers already upcast internally; this extends the same
         # treatment to attention and MLP.
-        # autocast to float32 is unsupported on CPU; skip fp32 upcast there
-        # (CPU uses full-precision intermediates internally).
+        # CPU autocast to float32 is unsupported, and Ascend fused operators
+        # require matching activation and weight dtypes. Preserve the existing
+        # stability fallback on the other accelerator backends.
         input_dtype = inputs_embeds.dtype
-        use_fp32 = input_dtype == torch.float16 and inputs_embeds.device.type != "cpu"
+        use_fp32 = input_dtype == torch.float16 and inputs_embeds.device.type not in ("cpu", "npu")
         if use_fp32:
             inputs_embeds = inputs_embeds.float()
         hidden_states = inputs_embeds
@@ -1035,6 +1036,9 @@ class CodePredictorWrapper(nn.Module):
         model_fwd = self._compiled_model_fwd
         lm_heads = self._lm_heads_list
         codec_embeds = self._codec_embeds_list
+        # torch-npu cannot replay an inner NPUGraph while the outer talker_mtp
+        # graph is being captured. Capture state is constant for this forward.
+        is_npu_capturing = current_omni_platform.is_npu() and torch.npu.is_current_stream_capturing()
 
         # Zero the padded region of the buffer
         proj_buf[:padded_bsz].zero_()
@@ -1085,7 +1089,9 @@ class CodePredictorWrapper(nn.Module):
             # Use captured device graph if available, otherwise call compiled fn.
             device_graph_entry = self._device_graphs.get(graph_key)
 
-            if device_graph_entry is not None:
+            # Let the outer graph record the regular forward during capture;
+            # normal inference still uses the inner graph replay fast path.
+            if device_graph_entry is not None and not is_npu_capturing:
                 device_graph_entry[0].replay()
                 hidden_out = device_graph_entry[1]
             else: