From 731e36fb5786e61fd569ba1979c6e3fdfe0db6ea Mon Sep 17 00:00:00 2001 From: yadong <76457085+yadongtan@users.noreply.github.com> Date: Tue, 15 Sep 2026 02:43:02 +0800 Subject: [PATCH] [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 Co-authored-by: tanyadong <2209931449@qq.com> --- .../qwen3_tts/test_code_predictor_dtype.py | 108 ++++++++++++++++++ .../models/common/qwen3_code_predictor.py | 14 ++- 2 files changed, 118 insertions(+), 4 deletions(-) diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py index 438a5ec1d..1998a4dd4 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py +++ b/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.""" diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 60b4d4525..eff3305d5 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/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: