diff --git a/docs/getting_started/installation/npu/npu.inc.md b/docs/getting_started/installation/npu/npu.inc.md index bc2d3b60c..a76363703 100644 --- a/docs/getting_started/installation/npu/npu.inc.md +++ b/docs/getting_started/installation/npu/npu.inc.md @@ -68,18 +68,18 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886 You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.) ```bash -# Pin vLLM version to 0.17.0 +# Pin vLLM version to 0.18.0 cd /vllm-workspace/vllm git pull origin main git fetch origin --tags -git checkout v0.17.0 +git checkout v0.18.0 VLLM_TARGET_DEVICE=empty pip install -v -e . # Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time. cd /vllm-workspace/vllm-ascend git pull origin main git fetch origin --tags -git checkout v0.17.0 +git checkout 1e05c4908f31737bc4eef865a9f351d030a77c9d pip install -v -e . # Install vLLM-Omni from the latest main branch diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 22badb49e..3d9cb86ba 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -26,6 +26,7 @@ from vllm.v1.outputs import ( from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.structured_output.utils import apply_grammar_bitmask from vllm.v1.utils import record_function_or_nullcontext +from vllm.v1.worker import mamba_utils from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm_ascend.ascend_forward_context import set_ascend_forward_context @@ -237,6 +238,23 @@ class NPUARModelRunner(OmniNPUModelRunner): pad_attn = cudagraph_mode == CUDAGraphMode.FULL + # NOTE(Angazenn): According to https://github.com/vllm-project/vllm/pull/30877, + # there should be a corresponding 'postprocess_mamba'. However, it is called inside + # '_update_states_after_model_execute', which is not overridden in vLLM-Ascend. + # We simply utilize the implementation in vLLM. + if self.cache_config.mamba_cache_mode == "align": + mamba_utils.preprocess_mamba( + scheduler_output, + self.kv_cache_config, + self.cache_config, + self.mamba_state_idx, + self.input_batch, + self.requests, + self.compilation_config.static_forward_context, + self.model.get_mamba_state_copy_func(), + self._get_mamba_copy_bufs(), + ) + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices @@ -339,7 +357,7 @@ class NPUARModelRunner(OmniNPUModelRunner): skip_compiled=has_encoder_input, ), self.maybe_get_kv_connector_output( - scheduler_output, clear_metadata=clear_kv_metadata + scheduler_output, defer_finalize=not clear_kv_metadata ) as kv_connector_output, ): hidden_states = self._model_forward( @@ -606,6 +624,34 @@ class NPUARModelRunner(OmniNPUModelRunner): hidden_states, multimodal_outputs, num_scheduled_tokens_np, scheduler_output ) + # Pre-copy multimodal tensors to CPU once (not per-request) to avoid + # redundant D2H transfers when gpu_resident_buffer_keys keeps them on GPU. + mm_cpu: dict[str, object] = {} + if isinstance(multimodal_outputs, dict) and multimodal_outputs: + for k, v in multimodal_outputs.items(): + try: + if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]: + mm_cpu[k] = v.detach().to("cpu").contiguous() + elif isinstance(v, dict): + sub_dict: dict[str, torch.Tensor] = {} + for sk, sv in v.items(): + if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]: + sub_dict[str(sk)] = sv.detach().to("cpu").contiguous() + if sub_dict: + mm_cpu[k] = sub_dict + elif isinstance(v, list): + if len(v) == 0: + continue + cpu_list = [] + for elem in v: + if isinstance(elem, torch.Tensor): + cpu_list.append(elem.detach().to("cpu").contiguous()) + else: + cpu_list.append(elem) + mm_cpu[k] = cpu_list + except Exception as e: + logger.error(f"Error in merge multimodal outputs: {e}") + pooler_output: list[dict[str, object]] = [] for rid in req_ids_output_copy: idx = req_id_to_index_output_copy[rid] @@ -614,28 +660,26 @@ class NPUARModelRunner(OmniNPUModelRunner): end = start + sched hidden_slice = hidden_states_cpu[start:end] payload: dict[str, object] = {"hidden": hidden_slice} - if isinstance(multimodal_outputs, dict) and multimodal_outputs: + if mm_cpu: mm_payload: dict[str, object] = {} - for k, v in multimodal_outputs.items(): - try: - if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]: - mm_payload[k] = v.detach().to("cpu")[start:end].contiguous() - elif isinstance(v, dict): - sub_dict: dict[str, torch.Tensor] = {} - for sk, sv in v.items(): - if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]: - sub_dict[str(sk)] = sv.detach().to("cpu")[start:end].contiguous() - if sub_dict: - mm_payload[k] = sub_dict - elif isinstance(v, list): - element = v[0] - if isinstance(element, torch.Tensor): - element = element.detach().to("cpu").contiguous() - mm_payload[k] = element - except Exception as e: - logger.error(f"Error in merge multimodal outputs: {e}") - if mm_payload: - payload.update(mm_payload) + for k, v in mm_cpu.items(): + if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]: + mm_payload[k] = v[start:end].contiguous() + elif isinstance(v, dict): + mm_payload[k] = {sk: sv[start:end].contiguous() for sk, sv in v.items()} + elif isinstance(v, list): + element = v[idx] if idx < len(v) else v[0] + # Clone tensors to avoid cross-request aliasing + if isinstance(element, torch.Tensor): + element = element.clone() + mm_payload[k] = element + elif isinstance(v, torch.Tensor): + # List-derived tensor payloads are request-invariant; clone to + # avoid accidental cross-request aliasing on downstream mutation. + mm_payload[k] = v.clone() + else: + mm_payload[k] = v + payload.update(mm_payload) pooler_output.append(payload) model_runner_output = OmniModelRunnerOutput( diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index 0296c042d..48a998280 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -20,6 +20,7 @@ from vllm.utils.math_utils import cdiv from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, make_empty_encoder_model_runner_output from vllm.v1.utils import record_function_or_nullcontext +from vllm.v1.worker import mamba_utils from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs @@ -195,6 +196,23 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): pad_attn = cudagraph_mode == CUDAGraphMode.FULL + # NOTE(Angazenn): According to https://github.com/vllm-project/vllm/pull/30877, + # there should be a corresponding 'postprocess_mamba'. However, it is called inside + # '_update_states_after_model_execute', which is not overridden in vLLM-Ascend. + # We simply utilize the implementation in vLLM. + if self.cache_config.mamba_cache_mode == "align": + mamba_utils.preprocess_mamba( + scheduler_output, + self.kv_cache_config, + self.cache_config, + self.mamba_state_idx, + self.input_batch, + self.requests, + self.compilation_config.static_forward_context, + self.model.get_mamba_state_copy_func(), + self._get_mamba_copy_bufs(), + ) + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices @@ -300,7 +318,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): skip_compiled=has_encoder_input, ), self.maybe_get_kv_connector_output( - scheduler_output, clear_metadata=clear_kv_metadata + scheduler_output, defer_finalize=not clear_kv_metadata ) as kv_connector_output, ): # -------------------------------------- Omni-new ------------------------------------------------- @@ -503,6 +521,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): remove_lora: bool = True, is_graph_capturing: bool = False, num_active_loras: int = 0, + profile_seq_lens: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: # only support eager mode and piecewise graph now assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes() @@ -609,11 +628,14 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): # seq_lens. We use this seq_len only when capturing graph, and still use max_query_len # in inference. This will be removed once npu_fused_infer_attention_score # outperforms _npu_paged_attention on all cases. - seq_lens = ( - SEQ_LEN_WITH_MAX_PA_WORKSPACE - if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config) - else max_query_len - ) # type: ignore[assignment] + if profile_seq_lens is not None: + seq_lens = profile_seq_lens + else: + seq_lens = ( + SEQ_LEN_WITH_MAX_PA_WORKSPACE + if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config) + else max_query_len + ) # type: ignore[assignment] self.seq_lens.np[:num_reqs_padded] = seq_lens self.seq_lens.np[num_reqs_padded:] = 0 self.seq_lens.copy_to_gpu() diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index 9ff5f720e..8ef39adfa 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -35,9 +35,11 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): enable_sp(self.vllm_config) # TODO move this model specific logic to a separate class # TTS model IS the talker (no .talker sub-attr); use getattr to support both Omni and TTS. + self.has_talker_mtp = False talker_mtp = getattr(self.model, "talker_mtp", None) if talker_mtp is not None: self.talker_mtp = talker_mtp # type: ignore[assignment] + self.has_talker_mtp = True cudagraph_mode = self.compilation_config.cudagraph_mode assert cudagraph_mode is not None # Only wrap talker_mtp in CUDAGraphWrapper for Omni models that @@ -73,6 +75,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): remove_lora: bool = True, is_graph_capturing: bool = False, num_active_loras: int = 0, + profile_seq_lens: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: # only support eager mode and piecewise graph now assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes() @@ -179,11 +182,14 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): # seq_lens. We use this seq_len only when capturing graph, and still use max_query_len # in inference. This will be removed once npu_fused_infer_attention_score # outperforms _npu_paged_attention on all cases. - seq_lens = ( - SEQ_LEN_WITH_MAX_PA_WORKSPACE - if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config) - else max_query_len - ) # type: ignore[assignment] + if profile_seq_lens is not None: + seq_lens = profile_seq_lens + else: + seq_lens = ( + SEQ_LEN_WITH_MAX_PA_WORKSPACE + if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config) + else max_query_len + ) # type: ignore[assignment] self.seq_lens.np[:num_reqs_padded] = seq_lens self.seq_lens.np[num_reqs_padded:] = 0 self.seq_lens.copy_to_gpu() @@ -283,7 +289,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): model_instance=self.model, ): # ---------------------------------------Omni-new---------------------------------------------- - if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"): + if getattr(self.model, "talker", None) is not None and self.has_talker_mtp: num_tokens_padded_talker_mtp = num_tokens_padded if num_tokens_padded_talker_mtp == self.max_num_tokens: num_tokens_padded_talker_mtp = self.talker_mtp_input_ids.gpu.shape[0] @@ -362,7 +368,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): # Omni-specific: wrap output if needed if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"): - model_output = self.model.make_omni_output(model_output, **model_kwargs_extra) + model_output = self.model.make_omni_output(model_output, **model_kwargs, **model_kwargs_extra) # Omni-specific: cache model output for later sample_tokens self._omni_last_model_output = model_output @@ -415,12 +421,14 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): None, self.vllm_config, aclgraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc ): req_embeds, code_predictor_codes = self.talker_mtp(req_input_ids, req_embeds, last_talker_hidden, text_step) - # update the inputs_embeds and code_predictor_codes - code_predictor_codes_cpu = code_predictor_codes.detach().to("cpu").contiguous() + # code_predictor_codes stays on GPU here; _update_intermediate_buffer + # keeps it device-resident when the key is in gpu_resident_buffer_keys. + # D2H is deferred to sample_tokens where hidden_states.to("cpu") already + # syncs the stream, avoiding a per-step cudaStreamSynchronize. out_key = getattr(self.model, "talker_mtp_output_key", "code_predictor_codes") for idx, req_id in enumerate(decode_req_ids): req_index = self.input_batch.req_ids.index(req_id) start_offset = int(self.query_start_loc.cpu[req_index]) inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1] - update_dict = {out_key: code_predictor_codes_cpu[idx : idx + 1]} + update_dict = {out_key: code_predictor_codes[idx : idx + 1]} self._merge_additional_information_update(req_id, update_dict)