From fe3d9a6a81e32a17ec9fb7c2ce215b6f19c0c8cf Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Fri, 26 Dec 2025 14:10:23 +0800 Subject: [PATCH] [NPU][Upgrade] Upgrade to v0.12.0 (#447) Signed-off-by: gcanlin Co-authored-by: Hongsheng Liu --- vllm_omni/patch.py | 63 + vllm_omni/worker/npu/npu_ar_model_runner.py | 1326 ++++++++++------- .../worker/npu/npu_generation_model_runner.py | 675 ++++++--- vllm_omni/worker/npu/npu_model_runner.py | 380 +++-- 4 files changed, 1579 insertions(+), 865 deletions(-) diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 687ff5186..9ecb5eeb0 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -14,6 +14,7 @@ from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs, OmniEn from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding from vllm_omni.request import OmniRequest +from vllm_omni.utils import is_npu for module_name, module in sys.modules.items(): # only do patch on module of vllm, pass others @@ -31,3 +32,65 @@ for module_name, module in sys.modules.items(): module.Request = OmniRequest if hasattr(module, "EngineCoreRequest") and module.EngineCoreRequest == _OriginalEngineCoreRequest: module.EngineCoreRequest = OmniEngineCoreRequest + + +# Patch for vllm-ascend prefetch functions bug fix +# Issue: The original functions access forward_context attributes like +# prefetch_mlp_gate_up_proj, prefetch_mlp_down_proj, layer_idx without checking +# if they exist, which causes AttributeError when prefetch_mlp_enabled is not set. +# TODO: Remove this patch after upgrading to vllm-ascend v0.13.0 or later. +# This issue has been fixed in https://github.com/vllm-project/vllm-ascend/pull/5035 +if is_npu(): + import torch + import torch.nn as nn + from vllm.model_executor.models.qwen2_5_omni_thinker import Qwen2_5_VLImageInputs, Qwen2_5_VLVideoInputs + from vllm_ascend.ascend_forward_context import set_ascend_forward_context + + from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( + Qwen2_5OmniThinkerForConditionalGeneration, + ) + + class AscendQwen2_5OmniThinkerForConditionalGeneration(nn.Module): + def _process_image_input(self, image_input: Qwen2_5_VLImageInputs) -> tuple[torch.Tensor, ...]: + if image_input["type"] == "image_embeds": + return image_input["image_embeds"].type(self.visual.dtype) + + grid_thw = image_input["image_grid_thw"] + assert grid_thw.ndim == 2 + + pixel_values = image_input["pixel_values"].type(self.visual.dtype) + with set_ascend_forward_context(None, self.vllm_config): + image_embeds = self.visual(pixel_values, grid_thw=grid_thw) + # Split concatenated embeddings for each image item. + merge_size = self.visual.spatial_merge_size + sizes = grid_thw.prod(-1) // merge_size // merge_size + + return image_embeds.split(sizes.tolist()) + + def _process_video_input( + self, + video_input: Qwen2_5_VLVideoInputs, + video_hashes: list[str] | None = None, + cached_video_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + if video_input["type"] == "video_embeds": + return video_input["video_embeds"].type(self.visual.dtype) + + grid_thw = video_input["video_grid_thw"] + assert grid_thw.ndim == 2 + + pixel_values_videos = video_input["pixel_values_videos"].type(self.visual.dtype) + with set_ascend_forward_context(None, self.vllm_config): + video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw) + # Split concatenated embeddings for each video item. + merge_size = self.visual.spatial_merge_size + sizes = grid_thw.prod(-1) // merge_size // merge_size + + return video_embeds.split(sizes.tolist()) + + Qwen2_5OmniThinkerForConditionalGeneration._process_image_input = ( + AscendQwen2_5OmniThinkerForConditionalGeneration._process_image_input + ) + Qwen2_5OmniThinkerForConditionalGeneration._process_video_input = ( + AscendQwen2_5OmniThinkerForConditionalGeneration._process_video_input + ) diff --git a/vllm_omni/worker/npu/npu_ar_model_runner.py b/vllm_omni/worker/npu/npu_ar_model_runner.py index bff15f2ab..ac677e3cd 100644 --- a/vllm_omni/worker/npu/npu_ar_model_runner.py +++ b/vllm_omni/worker/npu/npu_ar_model_runner.py @@ -4,55 +4,89 @@ from __future__ import annotations import math -from typing import Any +from typing import Any, NamedTuple import numpy as np import torch import torch.nn as nn from vllm.config import CUDAGraphMode -from vllm.distributed import tensor_model_parallel_all_gather +from vllm.distributed import get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather +from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group -from vllm.distributed.parallel_state import get_pp_group, get_tp_group -from vllm.forward_context import BatchDescriptor, get_forward_context +from vllm.distributed.parallel_state import get_pcp_group, get_pp_group, get_tp_group +from vllm.forward_context import get_forward_context from vllm.logger import logger from vllm.sequence import IntermediateTensors from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder -from vllm.v1.core.sched.output import SchedulerOutput - -# yapf conflicts with isort for this block -# yapf: disable -from vllm.v1.kv_cache_interface import ( - EncoderOnlyAttentionSpec, -) - -# yapf: enable +from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput +from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec from vllm.v1.outputs import ( EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, ModelRunnerOutput, + make_empty_encoder_model_runner_output, ) from vllm.v1.spec_decode.metadata import SpecDecodeMetadata +from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorOutput from vllm_ascend.ascend_forward_context import set_ascend_forward_context from vllm_ascend.attention.attention_v1 import AscendAttentionState from vllm_ascend.attention.utils import AscendCommonAttentionMetadata -from vllm_ascend.compilation.acl_graph import update_attn_params, update_mla_attn_params -from vllm_ascend.spec_decode.interface import SpecDcodeType -from vllm_ascend.utils import ( - ProfileExecuteDuration, - enable_sp, - lmhead_tp_enable, -) -from vllm_ascend.worker.model_runner_v1 import AsyncNPUModelRunnerOutput -from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload +# yapf conflicts with isort for this block +# yapf: disable +from vllm_ascend.compilation.acl_graph import ( + update_attn_dcp_pcp_params, + update_attn_params, + update_mla_attn_dcp_pcp_params, + update_mla_attn_params, +) +from vllm_ascend.patch.worker.patch_module import patch_torch_npu_argsort +from vllm_ascend.spec_decode.interface import SpecDcodeType +from vllm_ascend.utils import ProfileExecuteDuration, enable_sp, lmhead_tp_enable + from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.worker.npu.npu_model_runner import OmniNPUModelRunner +class ExecuteModelState(NamedTuple): + """Ephemeral cached state transferred between execute_model() and + sample_tokens(), after execute_model() returns None.""" + + scheduler_output: SchedulerOutput + logits: torch.Tensor + spec_decode_metadata: SpecDecodeMetadata | None + hidden_states: torch.Tensor + sample_hidden_states: torch.Tensor + aux_hidden_states: list[torch.Tensor] | None + kv_connector_output: KVConnectorOutput | None + attn_metadata: dict[str, Any] + positions: torch.Tensor + multimodal_outputs: Any + class NPUARModelRunner(OmniNPUModelRunner): """Autoregressive NPU model runner that returns hidden states per request.""" + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32) + # each model stage has their own hidden size + self.hidden_size = self.model_config.hf_text_config.hidden_size + self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False) + + def _make_buffer(self, *size, dtype, numpy=True): + # Prevent ray from pinning the buffer due to large size + from vllm_omni.distributed.ray_utils.utils import ( + calculate_total_bytes, + maybe_disable_pin_memory_for_ray, + ) + + total_bytes = calculate_total_bytes(size, dtype) + + # Use the context manager to temporarily disable pinning if needed + with maybe_disable_pin_memory_for_ray(self, total_bytes): + return super()._make_buffer(*size, dtype=dtype, numpy=numpy) + def _prepare_inputs( self, scheduler_output: SchedulerOutput, @@ -69,8 +103,7 @@ class NPUARModelRunner(OmniNPUModelRunner): torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, - int, - dict[str, dict] | None, + int ]: total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens assert total_num_scheduled_tokens > 0 @@ -85,45 +118,74 @@ class NPUARModelRunner(OmniNPUModelRunner): req_ids = self.input_batch.req_ids tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids] num_scheduled_tokens = np.array(tokens, dtype=np.int32) - max_num_scheduled_tokens = num_scheduled_tokens.max() - num_valid_tokens = np.array( - [ - num_tokens - len(scheduler_output.scheduled_spec_decode_tokens.get(i, [])) - for num_tokens, i in zip(tokens, req_ids) - ], - dtype=np.int32, + + req_indices = np.repeat(self.arange_np[:num_reqs], + num_scheduled_tokens) + _, arange = self._get_cumsum_and_arange(num_scheduled_tokens) + positions_np = np.add( + self.input_batch.num_computed_tokens_cpu[req_indices], + arange, ) - if self.use_aclgraph and total_num_scheduled_tokens <= self.aclgraph_batch_sizes[-1]: + self.input_batch.block_table.compute_slot_mapping( + req_indices, positions_np) + self.input_batch.block_table.commit_slot_mapping( + total_num_scheduled_tokens) + if self.pcp_size > 1: + if not self.vllm_config.model_config.use_mla: + self.generate_kv_idx(scheduler_output) + tokens, position_pcp, pcp_unpad_mask = self._update_tokens_for_pcp( + tokens) + num_scheduled_tokens = np.array(tokens, dtype=np.int32) + total_num_scheduled_tokens = sum(num_scheduled_tokens[:num_reqs]) + else: + position_pcp, pcp_unpad_mask = None, None + self.num_pcp_pads = self.num_pcp_pads[:num_reqs] + + total_num_pcp_pads = sum(self.num_pcp_pads) + max_num_scheduled_tokens = max(tokens) + num_valid_tokens = np.array([ + num_tokens - + len(scheduler_output.scheduled_spec_decode_tokens.get(i, [])) + for num_tokens, i in zip(tokens, req_ids) + ], + dtype=np.int32) + + if (self.use_aclgraph and total_num_scheduled_tokens + <= self.cudagraph_batch_sizes[-1]): # Add padding to the batch size. - num_input_tokens = self.vllm_config.pad_for_cudagraph(total_num_scheduled_tokens) + num_input_tokens = self.vllm_config.pad_for_cudagraph( + total_num_scheduled_tokens) elif self.use_aclgraph and enable_sp(self.vllm_config): # When using aclgraph, if total_num_scheduled_tokens exceeds the maximum graph size, # the model will fall back to running its FX graph in eager mode. # In this case, when sequence parallelism is enabled, we need to pad tokens to align # with tp_size because pad_size cannot be captured by the FX graph tp_size = self.vllm_config.parallel_config.tensor_parallel_size - num_input_tokens = math.ceil(total_num_scheduled_tokens / tp_size) * tp_size + num_input_tokens = math.ceil( + total_num_scheduled_tokens / tp_size) * tp_size else: # Eager mode. num_input_tokens = total_num_scheduled_tokens # Get the attention state. - attn_state = self._build_attn_state(num_reqs, num_scheduled_tokens, num_valid_tokens) + attn_state = self._build_attn_state(num_reqs, num_scheduled_tokens, + num_valid_tokens) self.attn_state = attn_state # type: ignore # Determine if it's a splitfuse batch - with_prefill = attn_state not in [AscendAttentionState.DecodeOnly, AscendAttentionState.SpecDecoding] + with_prefill = attn_state not in [ + AscendAttentionState.DecodeOnly, AscendAttentionState.SpecDecoding + ] self.query_lens = torch.from_numpy(num_scheduled_tokens) - enable_dbo = self._check_dbo_is_valid(self.query_lens.tolist(), attn_state, total_num_scheduled_tokens) # Get info across DP ranks. # NOTE: maybe_padded_num_tokens is only used when using TorchAir with DP, # Otherwise, it's just max_tokens_across_dp_cpu - (maybe_padded_num_tokens, num_tokens_across_dp, with_prefill, enable_dbo) = self._sync_metadata_across_dp( - num_input_tokens, with_prefill, enable_dbo - ) + (maybe_padded_num_tokens, num_tokens_across_dp, + with_prefill) = self._sync_metadata_across_dp(num_input_tokens, + with_prefill) # TODO: Now that num_input_tokens is basically identical with maybe_padded_num_tokens # We should consider removing maybe_padded_num_tokens later @@ -135,14 +197,21 @@ class NPUARModelRunner(OmniNPUModelRunner): # Get request indices. # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2] - req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens) + req_indices = np.repeat(self.arange_np[:num_reqs], + num_scheduled_tokens) # cu_num_tokens: [2, 5, 3] -> [2, 7, 10] # arange: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2] - cu_num_tokens, arange = self._get_cumsum_and_arange(num_scheduled_tokens) + cu_num_tokens, arange = self._get_cumsum_and_arange( + num_scheduled_tokens) - positions_np = self.positions_np[:total_num_scheduled_tokens] - np.add(self.input_batch.num_computed_tokens_cpu[req_indices], arange, out=positions_np) + if self.pcp_size > 1: + positions_np = self.positions.np[:total_num_scheduled_tokens] + np.add(self.input_batch.num_computed_tokens_cpu[req_indices], + position_pcp[:total_num_scheduled_tokens], + out=positions_np) + else: + self.positions.np[:total_num_scheduled_tokens] = positions_np # Calculate M-RoPE positions. # Only relevant for models using M-RoPE (e.g, Qwen2-VL) @@ -150,163 +219,259 @@ class NPUARModelRunner(OmniNPUModelRunner): self._calc_mrope_positions(scheduler_output) # Only relevant for models using M-RoPE (e.g, Qwen2-VL) - self.mrope_positions[:, :total_num_scheduled_tokens].copy_( - self.mrope_positions_cpu[:, :total_num_scheduled_tokens], non_blocking=True - ) + self.mrope_positions.gpu[:, :total_num_scheduled_tokens].copy_( + self.mrope_positions.cpu[:, :total_num_scheduled_tokens], + non_blocking=True) # Get token indices. # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2] # -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2] # where M is the max_model_len. - token_indices = positions_np + req_indices * self.input_batch.token_ids_cpu.shape[1] - + token_indices = (positions_np + + req_indices * self.input_batch.token_ids_cpu.shape[1]) + token_indices_tensor = torch.from_numpy(token_indices) # Prepare input_ids. # NOTE(woosuk): We use torch.index_select instead of np.take here # because torch.index_select is much faster than np.take for large # tensors. - torch.index_select( - self.input_batch.token_ids_cpu_tensor.flatten(), - 0, - torch.from_numpy(token_indices), - out=self.input_ids_cpu[:total_num_scheduled_tokens], - ) + torch.index_select(self.input_batch.token_ids_cpu_tensor.flatten(), + 0, + token_indices_tensor, + out=self.input_ids.cpu[:total_num_scheduled_tokens]) + if self.enable_prompt_embeds: + is_token_ids = self.input_batch.is_token_ids_tensor.flatten() + torch.index_select( + is_token_ids, + 0, + token_indices_tensor, + out=self.is_token_ids.cpu[:total_num_scheduled_tokens]) - # Prepare some information for building Attention-Metadata - # Compute and commit slot mapping - self.input_batch.block_table.compute_slot_mapping(req_indices, positions_np) - self.input_batch.block_table.commit_slot_mapping(total_num_scheduled_tokens) + # Because we did not pre-allocate a massive prompt_embeds CPU tensor on + # the InputBatch, we need to fill in the prompt embeds into the expected + # spots in the GpuModelRunner's pre-allocated prompt_embeds tensor. + if self.input_batch.req_prompt_embeds and (self.is_multimodal_model or + self.enable_prompt_embeds): + output_idx = 0 + for req_idx in range(num_reqs): + num_sched = num_scheduled_tokens[req_idx] - self.query_start_loc_np[0] = 0 - self.query_start_loc_np[1 : num_reqs + 1] = cu_num_tokens - self.query_start_loc[: num_reqs + 1].copy_(self.query_start_loc_cpu[: num_reqs + 1], non_blocking=True) + # Skip if this request doesn't have embeddings + if req_idx not in self.input_batch.req_prompt_embeds: + output_idx += num_sched + continue - self.seq_lens_np[:num_reqs] = self.input_batch.num_computed_tokens_cpu[:num_reqs] + num_scheduled_tokens - self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs], non_blocking=True) + # Skip if no tokens scheduled + if num_sched <= 0: + output_idx += num_sched + continue + + req_embeds = self.input_batch.req_prompt_embeds[req_idx] + start_pos = self.input_batch.num_computed_tokens_cpu[req_idx] + + # Skip if trying to read beyond available embeddings + if start_pos >= req_embeds.shape[0]: + output_idx += num_sched + continue + + # Copy available embeddings + end_pos = start_pos + num_sched + actual_end = min(end_pos, req_embeds.shape[0]) + actual_num_sched = actual_end - start_pos + + if actual_num_sched > 0: + self.inputs_embeds.cpu[output_idx:output_idx + + actual_num_sched].copy_( + req_embeds[start_pos:actual_end] + ) + + output_idx += num_sched + + self.query_start_loc.np[0] = 0 + self.query_start_loc.np[1:num_reqs + 1] = cu_num_tokens + self.query_start_loc.copy_to_gpu() + + self.seq_lens.np[:num_reqs] = ( + self.input_batch.num_computed_tokens_cpu[:num_reqs] + + num_scheduled_tokens) + self.seq_lens.copy_to_gpu() # Fill unused with -1. Needed for reshape_and_cache - self.query_start_loc[num_reqs + 1 :].fill_(-1) - self.seq_lens[num_reqs:].fill_(0) + self.query_start_loc.gpu[num_reqs + 1:].fill_(-1) + self.seq_lens.gpu[num_reqs:].fill_(0) self.query_lens = torch.from_numpy(num_scheduled_tokens) # Copy the tensors to the NPU. - self._prepare_input_ids(total_num_scheduled_tokens, cu_num_tokens) - self.positions_cpu[total_num_scheduled_tokens:num_input_tokens].zero_() - self.positions[:num_input_tokens].copy_(self.positions_cpu[:num_input_tokens], non_blocking=True) + self._prepare_input_ids(scheduler_output, total_num_scheduled_tokens, + cu_num_tokens) + self.positions.cpu[total_num_scheduled_tokens:num_input_tokens].zero_() + self.positions.copy_to_gpu() - # Make Attention metadata - positions_cpu = self.positions_cpu[:num_input_tokens] - positions = self.positions[:num_input_tokens] - seq_lens_cpu = self.seq_lens_cpu[:num_reqs] - attn_state = self._build_attn_state(num_reqs, num_scheduled_tokens, num_valid_tokens) - self.attn_mask = self._make_attention_mask(seq_lens=seq_lens_cpu, position=positions_cpu, attn_state=attn_state) + attn_state = self._build_attn_state(num_reqs, num_scheduled_tokens, + num_valid_tokens) + self.attn_mask = self._make_attention_mask(attn_state) self.attn_state = attn_state # type: ignore self.with_prefill = with_prefill self.num_tokens_across_dp = num_tokens_across_dp - self._update_graph_pad_size(with_prefill, maybe_padded_num_tokens) attn_metadata: dict[str, Any] = {} - # Omni-new: per_req_additional_information - per_req_additional_information: dict[str, dict] | None = None + # Record the index of requests that should not be sampled, + # so that we could clear the sampled tokens before returning + num_tokens = [ + self.requests[r].num_tokens for r in self.input_batch.req_ids + ] + num_tokens_np = np.array(num_tokens, dtype=np.int32) + num_reqs = self.input_batch.num_reqs + if self.pcp_size > 1: + # while pcp > 1, we need the original num_scheduled_tokens before split + # to calculate discard_requests_mask + tokens_original = [ + scheduler_output.num_scheduled_tokens[i] for i in req_ids + ] + original_seq_lens_np = ( + self.input_batch.num_computed_tokens_cpu[:num_reqs] + + np.array(tokens_original, dtype=np.int32)) + discard_requests_mask = original_seq_lens_np < num_tokens_np + else: + discard_requests_mask = self.seq_lens.np[:num_reqs] < num_tokens_np + + discard_request_indices = np.nonzero(discard_requests_mask)[0] + self.num_discarded_requests = len(discard_request_indices) + self.discard_request_indices.np[:self.num_discarded_requests] = ( + discard_request_indices) + self.discard_request_indices.copy_to_gpu(self.num_discarded_requests) # _prepare_inputs may reorder the batch, so we must gather # multi-modal outputs after that to ensure the correct order if self.is_multimodal_model: - # Run the multimodal encoder if any. - self._execute_mm_encoder(scheduler_output) - mm_embeds = self._gather_mm_embeddings(scheduler_output) + with self.maybe_get_ec_connector_output( + scheduler_output, + encoder_cache=self.encoder_cache, + ): + # Run the multimodal encoder if any. + self._execute_mm_encoder(scheduler_output) + + # NOTE(woosuk): To unify token ids and soft tokens (vision + # embeddings), we always use embeddings (rather than token ids) + # as input to the multimodal model, even when the input is text. + input_ids = self.input_ids.gpu[:total_num_scheduled_tokens] + mm_embeds, is_mm_embed = self._gather_mm_embeddings( + scheduler_output) + + inputs_embeds = self.model.embed_input_ids( + input_ids, + multimodal_embeddings=mm_embeds, + is_multimodal=is_mm_embed, + ) - # NOTE(woosuk): To unify token ids and soft tokens (vision - # embeddings), we always use embeddings (rather than token ids) - # as input to the multimodal model, even when the input is text. - input_ids = self.input_ids[:total_num_scheduled_tokens] - if mm_embeds: - inputs_embeds = self.model.get_input_embeddings(input_ids, mm_embeds) - else: - inputs_embeds = self.model.get_input_embeddings(input_ids) # TODO(woosuk): Avoid the copy. Optimize. - self.inputs_embeds[:total_num_scheduled_tokens].copy_(inputs_embeds) - + self.inputs_embeds.gpu[:total_num_scheduled_tokens].copy_( + inputs_embeds) + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] # -------------------------------------- Omni-new ------------------------------------------------- - # Omni-new: Reset per-step additional information collector (deprecated concat path) - if hasattr(self, "_forward_additional_information"): - self._forward_additional_information = None - # Omni-new: per-request additional information for this step - per_req_additional_information = {} + input_ids = self.input_ids.gpu[:num_input_tokens] + # -------------------------------------- Omni-new ------------------------------------------------- + elif self.enable_prompt_embeds and get_pp_group().is_first_rank: + # Get the input embeddings for the tokens that are not input embeds, + # then put them into the appropriate positions. + # TODO(qthequartermasterman): Since even when prompt embeds are + # enabled, (a) not all requests will use prompt embeds, and (b) + # after the initial prompt is processed, the rest of the generated + # tokens will be token ids, it is not desirable to have the + # embedding layer outside of the acl graph all the time. The v0 + # engine avoids this by "double compiling" the acl graph, once + # with input_ids and again with inputs_embeds, for all num_tokens. + # If a batch only has token ids, then including the embedding layer + # in the acl graph will be more performant (like in the else case + # below). + token_ids_idx = self.is_token_ids.gpu[:total_num_scheduled_tokens] \ + .nonzero(as_tuple=False) \ + .squeeze(1) + # Some tokens ids may need to become embeds + if token_ids_idx.numel() > 0: + token_ids = self.input_ids.gpu[token_ids_idx] + tokens_to_embeds = self.model.embed_input_ids( + input_ids=token_ids) + self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds - # Omni-new: Overlay custom prompt_embeds per request for the prompt portion; - # collect additional_information (tensor/list) for prefill portion only - for req_index, req_id in enumerate(self.input_batch.req_ids): - req_state = self.requests[req_id] - pe_cpu = getattr(req_state, "prompt_embeds_cpu", None) - addi_cpu = getattr(req_state, "additional_information_cpu", None) - num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index]) - prompt_len = len(req_state.prompt_token_ids) - prompt_remaining = max(0, prompt_len - num_computed_tokens) - sched_tokens = int(num_scheduled_tokens[req_index]) - overlay_len = min(sched_tokens, prompt_remaining) - if overlay_len <= 0: - continue - if pe_cpu is not None: - src = pe_cpu[num_computed_tokens : num_computed_tokens + overlay_len].to( - dtype=self.dtype, device=self.device, non_blocking=True - ) - start_offset = int(self.query_start_loc_cpu[req_index]) - self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src) - # Build per-request additional information (no cross-request concat) - if addi_cpu is not None and isinstance(addi_cpu, dict): - req_info: dict[str, object] = {} - for k, v in addi_cpu.items(): - if isinstance(v, torch.Tensor): - # For prefill tokens, pass only the scheduled slice; - # for decode or no scheduled tokens, pass whole tensor - if overlay_len > 0: - try: - seg = ( - v[num_computed_tokens : num_computed_tokens + overlay_len] - .detach() - .to("cpu") - .contiguous() - ) - except Exception: - seg = v.detach().to("cpu").contiguous() - req_info[k] = seg - else: - req_info[k] = v.detach().to("cpu").contiguous() - elif isinstance(v, list): - req_info[k] = v - else: - req_info[k] = v - per_req_additional_information[req_id] = req_info - # ------------------------------------------------------------------------------------------------ - - inputs_embeds = self.inputs_embeds[:num_input_tokens] - input_ids = self.input_ids[:num_input_tokens] + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] + input_ids = None else: # For text-only models, we use token ids as input. # While it is possible to use embeddings as input just like the # multimodal models, it is not desirable for performance since # then the embedding layer is not included in the ACL graph. - input_ids = self.input_ids[:num_input_tokens] + input_ids = self.input_ids.gpu[:num_input_tokens] inputs_embeds = None - positions = self.positions[:num_input_tokens] - input_ids, positions = self._update_input_ids_and_positions( - input_ids, positions, num_input_tokens, with_prefill, maybe_padded_num_tokens - ) + positions = self.positions.gpu[:num_input_tokens] + if self.uses_mrope: + positions = self.mrope_positions.gpu[:, :num_input_tokens] + # -------------------------------------- Omni-new ------------------------------------------------- + self._omni_num_scheduled_tokens_np = num_scheduled_tokens + + # Note: only prefill need collect additional_information for now. + # Decode don't need per_req_additional_information anymore. + if inputs_embeds is not None: + # Prefill: overlay prompt_embeds and collect additional_information + self._collect_additional_information_for_prefill(num_scheduled_tokens) + + if hasattr(self.model, "has_preprocess") and self.model.has_preprocess: + # Overlay custom prompt_embeds per request for the prompt portion; + # collect additional_information (tensor/list) for prefill portion only + for req_index, req_id in enumerate(self.input_batch.req_ids): + req_state = self.requests.get(req_id) + req_infos = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None + + start_offset = int(self.query_start_loc.cpu[req_index]) + sched_tokens = int(num_scheduled_tokens[req_index]) + s, e = start_offset, start_offset + sched_tokens + span_len = int(e) - int(s) + + # call the custom process function + req_input_ids, req_embeds, update_dict = self.model.preprocess( + input_ids=input_ids[s:e], input_embeds=inputs_embeds[s:e], **req_infos + ) + # TODO(Peiqi): the merge stage could move out from the critical path + self._merge_additional_information_update(req_id, update_dict) + + # update the inputs_embeds and input_ids + seg_len = min(span_len, req_embeds.shape[0]) + inputs_embeds[s : s + seg_len] = req_embeds[:seg_len] + if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len: + input_ids[s : s + seg_len] = req_input_ids + + # -------------------------------------- Omni-new ------------------------------------------------- + + + # type: ignore if get_pp_group().is_first_rank: intermediate_tensors = None else: assert intermediate_tensors is not None assert self.intermediate_tensors is not None + # If both flashcomm1 and pp are used simultaneously, + # the shape of the received data and the shape of the space to be copied to will not match, + # requiring a recalculation of the incoming data's shape. + tp_size = get_tensor_model_parallel_world_size() + num_input_tokens_with_flashcomm1 = num_input_tokens + if enable_sp(): + num_input_tokens_with_flashcomm1 = (num_input_tokens + + tp_size - 1) // tp_size for k, v in intermediate_tensors.items(): - self.intermediate_tensors[k][:num_input_tokens].copy_(v[:num_input_tokens], non_blocking=True) - intermediate_tensors = IntermediateTensors( - {k: v[:num_input_tokens] for k, v in self.intermediate_tensors.items()} - ) + self.intermediate_tensors[ + k][:num_input_tokens_with_flashcomm1].copy_( + v[:num_input_tokens_with_flashcomm1], + non_blocking=True) + intermediate_tensors = IntermediateTensors({ + k: + v[:num_input_tokens_with_flashcomm1] + for k, v in self.intermediate_tensors.items() + }) - use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + use_spec_decode = len( + scheduler_output.scheduled_spec_decode_tokens) > 0 if not use_spec_decode: # NOTE(woosuk): Due to chunked prefills, the batch may contain # partial requests. While we should not sample any token @@ -314,35 +479,71 @@ class NPUARModelRunner(OmniNPUModelRunner): # We will ignore the sampled tokens from the partial requests. # TODO: Support prompt logprobs. spec_decode_metadata = None - logits_indices = torch.from_numpy(cu_num_tokens - 1).to(self.device, non_blocking=True) + if self.pcp_size * self.dcp_size > 1: + logits_indices = torch.from_numpy( + cu_num_tokens + ) * self.pcp_size - self.num_pcp_pads[:num_reqs] - 1 + logits_indices = logits_indices.pin_memory().to( + self.device, non_blocking=True) + else: + logits_indices = self.query_start_loc.gpu[1:num_reqs + 1] - 1 else: # Get the number of draft tokens for each request. # Iterate over the dictionary rather than all requests since not all # requests have draft tokens. num_draft_tokens = np.zeros(num_reqs, dtype=np.int32) - for req_id, draft_token_ids in scheduler_output.scheduled_spec_decode_tokens.items(): + # For chunked prefills, use -1 as mask rather than 0, as guided + # decoding may rollback speculative tokens. + num_decode_draft_tokens = np.full(num_reqs, -1, dtype=np.int32) + for req_id, draft_token_ids in ( + scheduler_output.scheduled_spec_decode_tokens.items()): req_idx = self.input_batch.req_id_to_index[req_id] num_draft_tokens[req_idx] = len(draft_token_ids) + num_decode_draft_tokens[req_idx] = (len(draft_token_ids) if ( + self.input_batch.num_computed_tokens_cpu[req_idx] + >= self.input_batch.num_prompt_tokens[req_idx]) else -1) - spec_decode_metadata = self._calc_spec_decode_metadata(num_draft_tokens, cu_num_tokens) + spec_decode_metadata = self._calc_spec_decode_metadata( + num_draft_tokens, cu_num_tokens, self.num_pcp_pads[:num_reqs]) logits_indices = spec_decode_metadata.logits_indices - self.num_draft_tokens.np[:num_reqs] = num_draft_tokens - self.num_draft_tokens.np[num_reqs:].fill(0) - self.num_draft_tokens.copy_to_gpu() + + # For DECODE only cuda graph of some attention backends (e.g., GDN). + self.num_decode_draft_tokens.np[: + num_reqs] = num_decode_draft_tokens + self.num_decode_draft_tokens.np[num_reqs:].fill(-1) + self.num_decode_draft_tokens.copy_to_gpu() + # save logits_indices for pcp spec decode usage + self.logits_indices = logits_indices # Used in the below loop. # query_start_loc_cpu = self.query_start_loc.cpu[:num_reqs + 1] - num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs] - spec_decode_common_attn_metadata = None + num_computed_tokens_cpu = ( + self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs]) + self.spec_decode_common_attn_metadata = None if use_spec_decode and self.need_accepted_tokens: - self.num_accepted_tokens.np[:num_reqs] = self.input_batch.num_accepted_tokens_cpu[:num_reqs] + self.num_accepted_tokens.np[:num_reqs] = ( + self.input_batch.num_accepted_tokens_cpu[:num_reqs]) self.num_accepted_tokens.np[num_reqs:].fill(1) self.num_accepted_tokens.copy_to_gpu() + if self.speculative_config and self.pcp_size > 1: + self._generate_pcp_mtp_input( + num_reqs, scheduler_output.total_num_scheduled_tokens, + scheduler_output.num_scheduled_tokens) + + long_seq_metadata = self._generate_pcp_metadata( + total_num_scheduled_tokens) # Prepare the attention metadata for each KV cache group and make layers # in the same group share the same metadata. - for kv_cache_group_id, kv_cache_group_spec in enumerate(self.kv_cache_config.kv_cache_groups): - if isinstance(kv_cache_group_spec.kv_cache_spec, EncoderOnlyAttentionSpec): + for kv_cache_group_id, kv_cache_group_spec in enumerate( + self.kv_cache_config.kv_cache_groups): + # NOTE: This is strange, why did we use total_num_scheduled_tokens before? + slot_mapping_size = (total_num_scheduled_tokens + if self.pcp_size == 1 else + total_num_scheduled_tokens * self.pcp_size - + total_num_pcp_pads) + if isinstance(kv_cache_group_spec.kv_cache_spec, + EncoderOnlyAttentionSpec): # Encoder-only layers do not have KV cache, so we need to # create a dummy block table and slot mapping for them. blk_table_tensor = torch.zeros( @@ -351,361 +552,398 @@ class NPUARModelRunner(OmniNPUModelRunner): device=self.device, ) slot_mapping = torch.zeros( - (total_num_scheduled_tokens,), + (total_num_scheduled_tokens, ), dtype=torch.int64, device=self.device, ) else: blk_table = self.input_batch.block_table[kv_cache_group_id] blk_table_tensor = blk_table.get_device_tensor() - slot_mapping = blk_table.slot_mapping_cpu[:total_num_scheduled_tokens] - self.slot_mapping[:total_num_scheduled_tokens].copy_( - slot_mapping[:total_num_scheduled_tokens], - non_blocking=True, - ) - self.slot_mapping[total_num_scheduled_tokens:].fill_(0) + blk_table.slot_mapping.gpu[slot_mapping_size:].fill_(0) + if self.pcp_size > 1: + slot_mapping_for_pcp = blk_table.slot_mapping.gpu[: + long_seq_metadata + . + num_actual_tokens_pcp_padded] + slot_mapping_for_pcp[slot_mapping_size:].fill_(-1) + assert pcp_unpad_mask is not None + pcp_padded_slot_mapping = self.pcp_padded_slot_mapping[: + pcp_unpad_mask + . + shape[ + 0]] + pcp_padded_slot_mapping.fill_(-1) + pcp_padded_slot_mapping[ + pcp_unpad_mask] = slot_mapping_for_pcp[: + slot_mapping_size] + slot_mapping_for_pcp[:long_seq_metadata. + num_actual_tokens_pcp_padded] = pcp_padded_slot_mapping + blk_table.slot_mapping.gpu[:long_seq_metadata.num_actual_tokens_pcp_padded] = \ + slot_mapping_for_pcp + slot_mapping = blk_table.slot_mapping.gpu + + # NOTE: This is a temporary hack, now in GPUModelRunner, this prepare_inputs + # has been split to multiple parts, and there are 3 parts that is related to this + # `num_reqs`, we'll take `query_start_loc` as an example: + # 1. self.query_start_loc.np[1 : num_reqs + 1] = cu_num_tokens + # 2. get `num_reqs_padded`, this depends on dispatcher and which is why we have the + # following simplified `dispatch` logic here, we try to minimize the impact + # 3. query_start_loc = self.query_start_loc.gpu[: num_reqs_padded + 1] + uniform_decode = (max_num_scheduled_tokens == self.uniform_decode_query_len) \ + and (total_num_scheduled_tokens == max_num_scheduled_tokens * num_reqs) + + # TODO: We should make this official ASAP. Also note that if we pad here, + # the builders won’t need to add any extra padding. + if self.compilation_config.cudagraph_mode.decode_mode() == CUDAGraphMode.FULL and \ + uniform_decode: + num_reqs_padded = num_input_tokens // self.uniform_decode_query_len + pad_size = num_reqs_padded - num_reqs + if pad_size > 0: + last_query_loc = self.query_start_loc.gpu[num_reqs] + + steps = torch.arange(1, + pad_size + 1, + device=self.device, + dtype=self.query_start_loc.gpu.dtype) + fill_values = last_query_loc + ( + steps * self.uniform_decode_query_len) + + self.query_start_loc.gpu[num_reqs + 1:num_reqs_padded + + 1] = fill_values + # So we are trying to simulate the behavior of GPUModelRunner's + # prepare_inputs for uniform decode mode by padding query_start_loc + num_reqs = num_reqs_padded # Make AscendCommonAttentionMetadata common_attn_metadata = AscendCommonAttentionMetadata( - query_start_loc=self.query_start_loc_cpu[: num_reqs + 1], - query_start_loc_cpu=self.query_start_loc_cpu[: num_reqs + 1], - seq_lens_cpu=self.seq_lens_cpu, - seq_lens=self.seq_lens_cpu[:num_reqs], + query_start_loc=self.query_start_loc.gpu[:num_reqs + 1], + query_start_loc_cpu=self.query_start_loc.cpu[:num_reqs + 1], + seq_lens_cpu=self.seq_lens.cpu[:num_reqs], + seq_lens=self.seq_lens.gpu[:num_reqs], num_reqs=num_reqs, - num_actual_tokens=total_num_scheduled_tokens, + num_actual_tokens=slot_mapping_size, num_input_tokens=num_input_tokens, actual_seq_lengths_q=self.actual_seq_lengths_q, # TODO: change this to the right block table for linear attn block_table_tensor=blk_table_tensor[:num_reqs], - slot_mapping=self.slot_mapping, + slot_mapping=slot_mapping, num_computed_tokens_cpu=num_computed_tokens_cpu, - positions=self.positions, + positions=self.positions.gpu, attn_mask=self.attn_mask, spec_attn_mask=self.spec_attn_mask, attn_state=self.attn_state, - enable_dbo_across_dp=enable_dbo, is_only_prefill=bool(np.all(num_valid_tokens != 1)), max_query_len=max_num_scheduled_tokens, - graph_pad_size=self.graph_pad_size, decode_token_per_req=self.decode_token_per_req, - cos=self.cos, - sin=self.sin, + prefill_context_parallel_metadata=long_seq_metadata, ) - if self.speculative_config and spec_decode_common_attn_metadata is None: - spec_decode_common_attn_metadata = common_attn_metadata + if self.speculative_config and self.pcp_size > 1: + # For pcp + spec decode, we flatten block_table + # to avoid irregular spec_attn_mask shape, e.g., + # num_decode_req=2, num_prefill_req=3, num_speculative_tokens=1, + # ori block_table: # [d0, d1, p0, p1, p2] + # (num_reqs_d + num_reqs_p, max_num_blocks), + # flattened block_table: [d0, d0, d1, d1, p0, p1, p2] + # (num_reqs_d * decode_threshold + num_reqs_p, max_num_blocks), + ori_query_lens = self.query_start_loc_pcp_full.cpu[1:num_reqs+1] - \ + self.query_start_loc_pcp_full.cpu[:num_reqs] + num_prefill_reqs = (ori_query_lens + > self.decode_threshold).sum().item() + num_decode_reqs = num_reqs - num_prefill_reqs + num_decode_reqs_flatten = num_decode_reqs * self.decode_threshold + blk_table_tensor[ + num_decode_reqs_flatten:num_decode_reqs_flatten + + num_prefill_reqs].copy_( + blk_table_tensor[num_decode_reqs:num_decode_reqs + + num_prefill_reqs].clone()) + blk_table_tensor[:num_decode_reqs_flatten].copy_( + blk_table_tensor[:num_decode_reqs].repeat_interleave( + self.decode_threshold, dim=0)) + common_attn_metadata.block_table_tensor = \ + blk_table_tensor[:num_decode_reqs_flatten + num_prefill_reqs] + + if self.speculative_config and \ + self.spec_decode_common_attn_metadata is None: + self.spec_decode_common_attn_metadata = common_attn_metadata for attn_group in self.attn_groups[kv_cache_group_id]: common_prefix_len = 0 extra_attn_metadata_args = {} builder = attn_group.get_metadata_builder() - if isinstance(builder, GDNAttentionMetadataBuilder) or self.model_config.runner_type == "pooling": + if isinstance(builder, GDNAttentionMetadataBuilder): if use_spec_decode: + patch_torch_npu_argsort() extra_attn_metadata_args = dict( - num_accepted_tokens=self.num_accepted_tokens.gpu[:num_reqs], - num_draft_tokens=self.num_draft_tokens.gpu[:num_reqs], + num_accepted_tokens=self.num_accepted_tokens. + gpu[:num_reqs], + num_decode_draft_tokens_cpu=self. + num_decode_draft_tokens.cpu[:num_reqs], ) attn_metadata_i = builder.build( common_prefix_len=common_prefix_len, common_attn_metadata=common_attn_metadata, - **extra_attn_metadata_args, - ) + **extra_attn_metadata_args) + elif self.model_config.runner_type == "pooling": + attn_metadata_i = builder.build( + common_prefix_len=common_prefix_len, + common_attn_metadata=common_attn_metadata, + **extra_attn_metadata_args) else: attn_metadata_i = builder.build( common_prefix_len=common_prefix_len, common_attn_metadata=common_attn_metadata, model=self.get_model(), - **extra_attn_metadata_args, - ) + **extra_attn_metadata_args) for layer_name in attn_group.layer_names: attn_metadata[layer_name] = attn_metadata_i if lmhead_tp_enable(): - max_num_reqs_across_dp = maybe_padded_num_tokens if not with_prefill else self.max_num_reqs - logits_indices = nn.functional.pad(logits_indices, (0, max_num_reqs_across_dp - logits_indices.shape[0])) + max_num_reqs_across_dp = self.max_num_reqs * self.uniform_decode_query_len + logits_indices = nn.functional.pad( + logits_indices, + (0, max_num_reqs_across_dp - logits_indices.shape[0])) - return ( - attn_metadata, - positions, - num_scheduled_tokens, - num_input_tokens, - num_tokens_across_dp, - maybe_padded_num_tokens, - logits_indices, - spec_decode_metadata, - input_ids, - inputs_embeds, - intermediate_tensors, - max_num_scheduled_tokens, - per_req_additional_information, - ) + return (attn_metadata, positions, num_scheduled_tokens, + num_input_tokens, num_tokens_across_dp, + maybe_padded_num_tokens, logits_indices, spec_decode_metadata, + input_ids, inputs_embeds, intermediate_tensors, + max_num_scheduled_tokens) @torch.inference_mode() def execute_model( self, scheduler_output: SchedulerOutput, intermediate_tensors: IntermediateTensors | None = None, - ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: - with ProfileExecuteDuration().capture_async("prepare input"): - self._update_states(scheduler_output) + ) -> OmniModelRunnerOutput | IntermediateTensors | None: + if self.execute_model_state is not None: + raise RuntimeError("State error: sample_tokens() must be called " + "after execute_model() returns None.") + with ProfileExecuteDuration().capture_async("prepare input"): # -------------------------------------- Omni-new ------------------------------------------------- - # Omni-new: Decode per-request prompt_embeds / additional_hidden_states payloads - # (if present) into CPU tensors - try: - new_reqs = getattr(scheduler_output, "scheduled_new_reqs", []) - if new_reqs: - for nr in new_reqs: - req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None) - if req_id is None: - continue - # prompt_embeds - payload_pe = getattr(nr, "prompt_embeds", None) - if payload_pe is not None: - if isinstance(payload_pe, torch.Tensor): - pe_cpu = payload_pe.detach().to("cpu").contiguous() - elif isinstance(payload_pe, PromptEmbedsPayload): - dt = np.dtype(getattr(payload_pe, "dtype", "float32")) - arr = np.frombuffer(payload_pe.data, dtype=dt) - arr = arr.reshape(payload_pe.shape) - pe_cpu = torch.from_numpy(arr) - else: - pe_cpu = None - if pe_cpu is not None and req_id in self.requests: - setattr( - self.requests[req_id], - "prompt_embeds_cpu", - pe_cpu, - ) - # additional_information - payload_info = getattr(nr, "additional_information", None) - if payload_info is not None: - info_dict = {} - if isinstance(payload_info, dict): - # Already decoded - info_dict = payload_info - elif isinstance(payload_info, AdditionalInformationPayload): - for k, entry in payload_info.entries.items(): - if entry.tensor_data is not None: - dt = np.dtype(getattr(entry, "tensor_dtype", "float32")) - arr = np.frombuffer(entry.tensor_data, dtype=dt) - arr = arr.reshape(entry.tensor_shape) - info_dict[k] = torch.from_numpy(arr) - else: - info_dict[k] = entry.list_data - if info_dict and req_id in self.requests: - setattr( - self.requests[req_id], - "additional_information_cpu", - info_dict, - ) - except Exception as e: - logger.error(f"Error decoding prompt_embeds / additional_information: {e}") - pass + self._update_states(scheduler_output) + self._decode_and_store_request_payloads(scheduler_output) # ------------------------------------------------------------------------------------------------ + if has_ec_transfer() and get_ec_transfer().is_producer: + with self.maybe_get_ec_connector_output( + scheduler_output, + encoder_cache=self.encoder_cache, + ): + self._execute_mm_encoder(scheduler_output) + return make_empty_encoder_model_runner_output( + scheduler_output) + if not scheduler_output.total_num_scheduled_tokens: if not has_kv_transfer_group(): - logger.debug("skip this step for we receive the data from remote disaggregate prefill node") + logger.debug( + "skip this step for we receive the data from remote disaggregate prefill node" + ) # Return empty ModelRunnerOutput if there's no work to do. return EMPTY_MODEL_RUNNER_OUTPUT - return self.kv_connector_no_forward(scheduler_output) + return self.kv_connector_no_forward(scheduler_output, + self.vllm_config) if self.dynamic_eplb: self.eplb_updator.forward_before() - ( - attn_metadata, - positions, - num_scheduled_tokens_np, - num_input_tokens, - num_tokens_across_dp, - maybe_padded_num_tokens, - logits_indices, - spec_decode_metadata, - input_ids, - inputs_embeds, - intermediate_tensors, - max_query_len, - per_req_additional_information, - ) = self._prepare_inputs(scheduler_output, intermediate_tensors) + (attn_metadata, positions, num_scheduled_tokens_np, + num_input_tokens, num_tokens_across_dp, maybe_padded_num_tokens, + logits_indices, spec_decode_metadata, input_ids, inputs_embeds, + intermediate_tensors, + max_query_len) = (self._prepare_inputs(scheduler_output, + intermediate_tensors)) if self.dynamic_eplb: self.eplb_updator.take_update_info_from_eplb_process() - moe_comm_type = self._select_moe_comm_method(num_input_tokens, self.with_prefill) + moe_comm_type = self._select_moe_comm_method(num_input_tokens) + # prevent debugger is None + need_dump = self.dump_enable and self.debugger is not None + if need_dump: + assert self.debugger is not None + dbg_cfg = getattr(self.debugger, "config", None) + dump_level = str( + getattr(dbg_cfg, "level", + "L1")).upper() if dbg_cfg is not None else "L1" + if dump_level in ("L0", "MIX"): + self.debugger.start(model=self.model) + else: + self.debugger.start() uniform_decode = (max_query_len == self.uniform_decode_query_len) and ( - scheduler_output.total_num_scheduled_tokens == self.input_batch.num_reqs * max_query_len - ) - batch_descriptor = BatchDescriptor(num_tokens=num_input_tokens, uniform_decode=uniform_decode) - aclgraph_runtime_mode, batch_descriptor = self.aclgraph_dispatcher.dispatch(batch_descriptor) + scheduler_output.total_num_scheduled_tokens + == self.input_batch.num_reqs * max_query_len) + has_lora = len(self.input_batch.lora_id_to_lora_request) > 0 + aclgraph_runtime_mode, batch_descriptor = \ + self.cudagraph_dispatcher.dispatch( + num_tokens=num_input_tokens, uniform_decode=uniform_decode, has_lora=has_lora) # Run forward pass with ProfileExecuteDuration().capture_async("forward"): with set_ascend_forward_context( - attn_metadata, - self.vllm_config, - num_tokens=num_input_tokens, - num_tokens_across_dp=num_tokens_across_dp, - with_prefill=self.with_prefill, - reserved_mc2_mask=self.reserved_mc2_mask, - moe_comm_type=moe_comm_type, - aclgraph_runtime_mode=aclgraph_runtime_mode, - batch_descriptor=batch_descriptor, - num_actual_tokens=scheduler_output.total_num_scheduled_tokens, - prefetch_stream=self.prefetch_stream, - model_instance=self.model, - weight_prefetch_method=self.weight_prefetch_method, - ): + attn_metadata, + self.vllm_config, + num_tokens=num_input_tokens, + num_tokens_across_dp=num_tokens_across_dp, + with_prefill=self.with_prefill, + moe_comm_type=moe_comm_type, + aclgraph_runtime_mode=aclgraph_runtime_mode, + batch_descriptor=batch_descriptor, + num_actual_tokens=scheduler_output. + total_num_scheduled_tokens, + prefetch_stream=self.prefetch_stream, + model_instance=self.model, + weight_prefetch_method=self.weight_prefetch_method): self.maybe_setup_kv_connector(scheduler_output) - # -------------------------------------- Omni-new ------------------------------------------------- - model_kwargs_extra = {} - # Pass per-request additional information map for this step (no concat) - if per_req_additional_information: - model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information - # Always pass per-request runtime additional_information (persisted in request state) - try: - per_req_runtime_info = [] - for req_id in self.input_batch.req_ids: - req_state = self.requests.get(req_id) - info = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None - per_req_runtime_info.append(info if isinstance(info, dict) else {}) - model_kwargs_extra["runtime_additional_information"] = per_req_runtime_info - model_kwargs_extra["request_ids"] = self.input_batch.req_ids - # Pass each request's token span within the flattened sequence for this step, - # enabling the model to map decode/prefill by request - req_token_spans = [] - for req_index in range(len(self.input_batch.req_ids)): - start_offset = int(self.query_start_loc_cpu[req_index]) - sched_tokens = int(num_scheduled_tokens_np[req_index]) - req_token_spans.append((start_offset, start_offset + sched_tokens)) - model_kwargs_extra["request_token_spans"] = req_token_spans - except Exception: - pass - # ------------------------------------------------------------------------------------------------ hidden_states = self._generate_process_reqs_hidden_states( - attn_metadata, - self.with_prefill, - maybe_padded_num_tokens, - input_ids, - positions, - intermediate_tensors, - inputs_embeds, - model_kwargs_extra, - ) + maybe_padded_num_tokens, input_ids, positions, + intermediate_tensors, inputs_embeds) self.maybe_wait_for_kv_save() - finished_sending, finished_recving = self.get_finished_kv_transfer(scheduler_output) + finished_sending, finished_recving = self.get_finished_kv_transfer( + scheduler_output) aux_hidden_states = None if self.drafter and self.drafter.name == SpecDcodeType.EAGLE3: hidden_states, aux_hidden_states = hidden_states - kv_connector_output = KVConnectorOutput(finished_sending=finished_sending, finished_recving=finished_recving) + kv_connector_output = KVConnectorOutput( + finished_sending=finished_sending, + finished_recving=finished_recving) finished_sending = None finished_recving = None with ProfileExecuteDuration().capture_async("post process"): + # -------------------------------------- Omni-new ------------------------------------------------- + multimodal_outputs = hidden_states.multimodal_outputs + hidden_states = hidden_states.text_hidden_states + + if multimodal_outputs is not None: + keys_or_type = ( + list(multimodal_outputs.keys()) + if isinstance(multimodal_outputs, dict) + else type(multimodal_outputs) + ) + logger.debug(f"[AR] execute_model: multimodal_outputs keys = {keys_or_type}") + else: + logger.debug("[AR] execute_model: multimodal_outputs is None") + # -------------------------------------- Omni-new ------------------------------------------------- # Broadcast PP output for external_launcher (torchrun) # to make sure we are synced across pp ranks # TODO: Support overlapping mirco-batches # https://github.com/vllm-project/vllm/issues/18019 - - # -------------------------------------- Omni-new ------------------------------------------------- - hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) - # The model side may return per-request additional_information updates (model-agnostic channel). - # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order; - # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode. - try: - if isinstance(multimodal_outputs, dict) and ( - "additional_information_update" in multimodal_outputs - or "additional_information_update_by_req_id" in multimodal_outputs - ): - # Option A: list[dict] in batch order - updates_list = multimodal_outputs.get("additional_information_update") - if isinstance(updates_list, list): - for idx, upd in enumerate(updates_list): - if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids): - continue - req_id = self.input_batch.req_ids[idx] - self._merge_additional_information_update(req_id, upd) - # Option B: dict[str, dict] keyed by req_id - updates_map = multimodal_outputs.get("additional_information_update_by_req_id") - if isinstance(updates_map, dict): - for req_id, upd in updates_map.items(): - if not isinstance(upd, dict): - continue - if req_id not in self.requests: - continue - self._merge_additional_information_update(req_id, upd) - except Exception as e: - logger.error( - f"Error merging for requests:{self.input_batch.req_ids} additional \ - information update: {e}, with the multimodal_outputs as {multimodal_outputs}" - ) - # ------------------------------------------------------------------------------------------------ - broadcast_pp_output = ( - self.parallel_config.distributed_executor_backend == "external_launcher" - and len(get_pp_group().ranks) > 0 - ) + broadcast_pp_output = \ + self.parallel_config.distributed_executor_backend \ + == "external_launcher" and len(get_pp_group().ranks) > 0 if not get_pp_group().is_last_rank: # For mid-pipeline stages, return the hidden states. if not broadcast_pp_output: hidden_states.kv_connector_output = kv_connector_output + if need_dump: + assert self.debugger is not None + self.debugger.stop() + self.debugger.step() return hidden_states assert isinstance(hidden_states, IntermediateTensors) - get_pp_group().send_tensor_dict(hidden_states.tensors, all_gather_group=get_tp_group()) + get_pp_group().send_tensor_dict( + hidden_states.tensors, all_gather_group=get_tp_group()) logits = None else: if self.input_batch.pooling_params: - return self._pool( + pool_output = self._pool( hidden_states, scheduler_output.total_num_scheduled_tokens, - num_scheduled_tokens_np, - finished_sending, - finished_recving, - kv_connector_output, - ) + num_scheduled_tokens_np) + if need_dump: + assert self.debugger is not None + self.debugger.stop() + self.debugger.step() + return pool_output + # Sometimes, after the model is compiled through the AOT backend, + # the model output may become a list containing only one Tensor object. + if isinstance(hidden_states, list) and \ + len(hidden_states) == 1 and \ + isinstance(hidden_states[0], torch.Tensor): + hidden_states = hidden_states[0] sample_hidden_states = hidden_states[logits_indices] logits = self.model.compute_logits(sample_hidden_states) if broadcast_pp_output: - model_output_broadcast_data = ( - { - "logits": logits.contiguous(), - } - if logits is not None - else {} - ) - model_output_broadcast_data = get_pp_group().broadcast_tensor_dict( - model_output_broadcast_data, src=len(get_pp_group().ranks) - 1 - ) + model_output_broadcast_data = { + "logits": logits.contiguous(), + } if logits is not None else {} + model_output_broadcast_data = get_pp_group( + ).broadcast_tensor_dict(model_output_broadcast_data, + src=len(get_pp_group().ranks) - 1) assert model_output_broadcast_data is not None logits = model_output_broadcast_data["logits"] # Apply structured output bitmasks if present - if scheduler_output.grammar_bitmask is not None: - logits = self.apply_grammar_bitmask(scheduler_output, logits) + self.execute_model_state = ExecuteModelState( + scheduler_output, + logits, + spec_decode_metadata, + hidden_states, + sample_hidden_states, + aux_hidden_states, + kv_connector_output, + attn_metadata, + positions, + multimodal_outputs # Omni-new + ) + return None + @torch.inference_mode + def sample_tokens( + self, grammar_output: GrammarOutput | None + ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: + if self.execute_model_state is None: + # Nothing to do (PP non-final rank case), output isn't used. + return None # noqa + need_dump = self.dump_enable and self.debugger is not None + # Unpack ephemeral state. + ( + scheduler_output, + logits, + spec_decode_metadata, + hidden_states, + sample_hidden_states, + aux_hidden_states, # noqa + kv_connector_output, + attn_metadata, + positions, + multimodal_outputs, # Omni-new + ) = self.execute_model_state + # Clear ephemeral state. + self.execute_model_state = None + + # Apply structured output bitmasks if present. + if grammar_output is not None: + logits = self.apply_grammar_bitmask(scheduler_output, + grammar_output, logits) + + with ProfileExecuteDuration().capture_async("Sample"): # Sample the next token and get logprobs if needed. sampling_metadata = self.input_batch.sampling_metadata if spec_decode_metadata is None: if lmhead_tp_enable() and logits is not None: - logits = logits[: self.input_batch.num_reqs] + logits = logits[:self.input_batch.num_reqs] sampler_output = self.sampler( logits=logits, sampling_metadata=sampling_metadata, ) else: if lmhead_tp_enable() and logits is not None: - logits = logits[: len(spec_decode_metadata.logits_indices)] + logits = logits[:len(spec_decode_metadata.logits_indices)] # When indexing with a tensor (bonus_logits_indices), PyTorch # creates a new tensor with separate storage from the original # logits tensor. This means any in-place operations on bonus_logits # won't affect the original logits tensor. assert logits is not None - bonus_logits = logits[spec_decode_metadata.bonus_logits_indices] + bonus_logits = logits[ + spec_decode_metadata.bonus_logits_indices] sampler_output = self.sampler( logits=bonus_logits, sampling_metadata=sampling_metadata, @@ -715,7 +953,8 @@ class NPUARModelRunner(OmniNPUModelRunner): # Just like `bonus_logits`, `target_logits` is a new tensor with # separate storage from the original `logits` tensor. Therefore, # it is safe to update `target_logits` in place. - target_logits = logits[spec_decode_metadata.target_logits_indices] + target_logits = logits[ + spec_decode_metadata.target_logits_indices] output_token_ids = self.rejection_sampler( spec_decode_metadata, None, # draft_probs @@ -726,68 +965,68 @@ class NPUARModelRunner(OmniNPUModelRunner): sampler_output.sampled_token_ids = output_token_ids if self.need_accepted_tokens: self._update_states_after_model_execute(output_token_ids) - - discard_sampled_tokens_req_indices: list[int] = [] - # TODO(woosuk): The following loop can be slow since it iterates over - # the requests one by one. Optimize. - discard_sampled_tokens_req_indices = [] - for i, req_id in enumerate(self.input_batch.req_ids): - req_state = self.requests[req_id] - seq_len = req_state.num_computed_tokens + scheduler_output.num_scheduled_tokens[req_id] - if seq_len < req_state.num_tokens: - # Ignore the sampled token. - # Rewind the generator state as if the token was not sampled. - generator = self.input_batch.generators.get(i) - if generator is not None: - generator.set_offset(generator.get_offset() - 4) - discard_sampled_tokens_req_indices.append(i) + discard_sampled_tokens_req_indices = \ + self.discard_request_indices.np[:self.num_discarded_requests] + for i in discard_sampled_tokens_req_indices: + generator = self.input_batch.generators.get(int(i)) + if generator is not None: + generator.set_offset(generator.get_offset() - 4) # Copy some objects so they don't get modified after returning. # This is important when using async scheduling. req_ids_output_copy = self.input_batch.req_ids.copy() - req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy() + req_id_to_index_output_copy = \ + self.input_batch.req_id_to_index.copy() # NOTE: NPU -> CPU Sync happens here. # Move as many CPU operations as possible before this sync point. logprobs_tensors = sampler_output.logprobs_tensors - logprobs_lists = logprobs_tensors.tolists() if logprobs_tensors is not None else None + logprobs_lists = logprobs_tensors.tolists() \ + if logprobs_tensors is not None else None # Compute prompt logprobs if needed. prompt_logprobs_dict = self._get_prompt_logprobs_dict( - hidden_states[: scheduler_output.total_num_scheduled_tokens], + hidden_states[:scheduler_output.total_num_scheduled_tokens], scheduler_output, ) num_sampled_tokens = sampler_output.sampled_token_ids.shape[0] sampled_token_ids = sampler_output.sampled_token_ids + if not self.use_async_scheduling: # Get the valid generated tokens. max_gen_len = sampled_token_ids.shape[-1] if max_gen_len == 1: - # No spec decode tokens. + # No spec decode tokens. It's a tensor. valid_sampled_token_ids = sampled_token_ids.tolist() else: - # Includes spec decode tokens. - valid_sampled_token_ids = self.rejection_sampler.parse_output( + # Includes spec decode tokens. It's a numpy array + valid_sampled_token_ids, _ = self.rejection_sampler.parse_output( sampled_token_ids, self.input_batch.vocab_size, ) # Mask out the sampled tokens that should not be sampled. for i in discard_sampled_tokens_req_indices: - valid_sampled_token_ids[i].clear() + valid_sampled_token_ids[int(i)].clear() else: valid_sampled_token_ids = [] - invalid_req_indices = list(discard_sampled_tokens_req_indices) + invalid_req_indices = discard_sampled_tokens_req_indices.tolist( + ) invalid_req_indices_set = set(invalid_req_indices) - assert sampled_token_ids.shape[-1] == 1 + if self.num_spec_tokens <= 0: + assert sampled_token_ids.shape[-1] == 1 + # Cache the sampled tokens on the NPU and avoid CPU sync. + # These will be copied into input_ids in the next step + # when preparing inputs. + self.input_batch.prev_sampled_token_ids = sampled_token_ids - # Cache the sampled tokens on the NPU and avoid CPU sync. - # These will be copied into input_ids in the next step - # when preparing inputs. - self.input_batch.prev_sampled_token_ids = sampled_token_ids - self.input_batch.prev_sampled_token_ids_invalid_indices = invalid_req_indices_set + + self.input_batch.prev_sampled_token_ids_invalid_indices = \ + invalid_req_indices_set self.input_batch.prev_req_id_to_index = { - req_id: i for i, req_id in enumerate(self.input_batch.req_ids) if i not in invalid_req_indices_set + req_id: i + for i, req_id in enumerate(self.input_batch.req_ids) + if i not in invalid_req_indices_set } # Cache the sampled tokens in the model runner, so that the scheduler # doesn't need to send them back. @@ -796,7 +1035,8 @@ class NPUARModelRunner(OmniNPUModelRunner): # between the first-stage worker and the last-stage worker. for req_idx in range(num_sampled_tokens): if self.use_async_scheduling: - sampled_ids = [-1] * 1 if req_idx not in invalid_req_indices_set else None + sampled_ids = [-1] * 1 if \ + req_idx not in invalid_req_indices_set else None else: sampled_ids = valid_sampled_token_ids[req_idx] if not sampled_ids: @@ -807,73 +1047,94 @@ class NPUARModelRunner(OmniNPUModelRunner): assert end_idx <= self.model_config.max_model_len, ( "Sampled token IDs exceed the max model length. " f"Total number of tokens: {end_idx} > max_model_len: " - f"{self.model_config.max_model_len}" - ) + f"{self.model_config.max_model_len}") - self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = sampled_ids + self.input_batch.token_ids_cpu[req_idx, + start_idx:end_idx] = sampled_ids + self.input_batch.is_token_ids[req_idx, + start_idx:end_idx] = True self.input_batch.num_tokens_no_spec[req_idx] = end_idx self.input_batch.num_tokens[req_idx] = end_idx req_id = self.input_batch.req_ids[req_idx] req_state = self.requests[req_id] req_state.output_token_ids.extend(sampled_ids) + def propose_draft_token_ids(sampled_token_ids): + assert self.spec_decode_common_attn_metadata is not None + self._draft_token_ids = self.propose_draft_token_ids( + sampled_token_ids, + sampling_metadata, + scheduler_output, + spec_decode_metadata, + positions, + scheduler_output.total_num_scheduled_tokens, + hidden_states, + attn_metadata, + aux_hidden_states, + ) + + with ProfileExecuteDuration().capture_async("Draft"): if self.speculative_config: - self._draft_token_ids = self.propose_draft_token_ids( - valid_sampled_token_ids, - sampling_metadata, - scheduler_output, - spec_decode_metadata, - positions, - scheduler_output.total_num_scheduled_tokens, - hidden_states, - attn_metadata, - aux_hidden_states, - ) + use_padded_batch_for_eagle = self.speculative_config and \ + self.speculative_config.method == "mtp" and \ + not self.speculative_config.disable_padded_drafter_batch + if use_padded_batch_for_eagle: + # EAGLE speculative decoding can use the GPU sampled tokens + # as inputs, and does not need to wait for bookkeeping to finish. + propose_draft_token_ids(sampler_output.sampled_token_ids) + if self.speculative_config and not use_padded_batch_for_eagle: + # ngram and other speculative decoding methods use the sampled + # tokens on the CPU, so they are run after bookkeeping. + propose_draft_token_ids(valid_sampled_token_ids) if has_kv_transfer_group(): get_kv_transfer_group().clear_connector_metadata() # -------------------------------------- Omni-new ------------------------------------------------- - # Omni-new: Convert to per-request tensors on CPU hidden_states_cpu = hidden_states.detach().to("cpu").contiguous() - pooler_output: list[torch.Tensor | None] = [] - prev_logits_index = 0 - for rid, logits_index in zip(req_ids_output_copy, logits_indices): - # Base payload: hidden slice for this request in this iteration - hidden_slice = hidden_states_cpu[prev_logits_index : logits_index + 1] + num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None) + if num_scheduled_tokens_np is None: + req_ids = self.input_batch.req_ids + num_scheduled_tokens_np = np.array( + [scheduler_output.num_scheduled_tokens[rid] for rid in req_ids], + dtype=np.int32, + ) + + self._process_additional_information_updates(hidden_states, multimodal_outputs, num_scheduled_tokens_np) + + pooler_output: list[dict[str, object]] = [] + for rid in req_ids_output_copy: + idx = req_id_to_index_output_copy[rid] + start = int(self.query_start_loc.cpu[idx]) + sched = int(num_scheduled_tokens_np[idx]) + end = start + sched + hidden_slice = hidden_states_cpu[start:end] payload: dict[str, object] = {"hidden": hidden_slice} - # Merge multimodal_outputs if present if isinstance(multimodal_outputs, dict) and multimodal_outputs: mm_payload: dict[str, object] = {} for k, v in multimodal_outputs.items(): try: - # Case 1: tensor aligned on token dimension if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]: - mm_payload[k] = v.detach().to("cpu")[prev_logits_index : logits_index + 1].contiguous() - # Case 2: nested dict of tensors aligned on token dimension (e.g., selected_hidden_layers) + 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")[prev_logits_index : logits_index + 1].contiguous() - ) + sub_dict[str(sk)] = sv.detach().to("cpu")[start:end].contiguous() if sub_dict: mm_payload[k] = sub_dict elif isinstance(v, list): - element: torch.Tensor = v[0] - multimodal_outputs[k] = v[1:] if len(v) > 1 else v + element = v[0] + if isinstance(element, torch.Tensor): + element = element.detach().to("cpu").contiguous() mm_payload[k] = element except Exception as e: - # Best-effort; skip malformed entries logger.error(f"Error in merge multimodal outputs: {e}") if mm_payload: payload.update(mm_payload) - pooler_output.append(payload) # type: ignore[arg-type] - prev_logits_index = logits_index + 1 + pooler_output.append(payload) - # Omni-new - output = OmniModelRunnerOutput( + model_runner_output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, sampled_token_ids=valid_sampled_token_ids, @@ -882,23 +1143,37 @@ class NPUARModelRunner(OmniNPUModelRunner): pooler_output=(pooler_output if self.vllm_config.model_config.engine_output_type != "text" else None), kv_connector_output=kv_connector_output, ) - # ------------------------------------------------------------------------------------------------ + # -------------------------------------- Omni-new ------------------------------------------------- durations = ProfileExecuteDuration().pop_captured_sync() if durations: - dr_str = [f"[{tag}]:{duration:.2f}ms" for tag, duration in durations.items()] + dr_str = [ + f"[{tag}]:{duration:.2f}ms" + for tag, duration in durations.items() + ] captured_name = "Decode" if self.attn_state == AscendAttentionState.DecodeOnly else "Prefill" - logger.info("Profile execute duration [%s]:%s", captured_name, " ".join(dr_str)) + logger.info("Profile execute duration [%s]:%s", captured_name, + " ".join(dr_str)) if self.dynamic_eplb: self.eplb_updator.forward_end() if not self.use_async_scheduling: - return output + if need_dump: + assert self.debugger is not None + self.debugger.stop() + self.debugger.step() + return model_runner_output - return AsyncNPUModelRunnerOutput( - model_runner_output=output, + if need_dump: + assert self.debugger is not None + self.debugger.stop() + self.debugger.step() + return AsyncGPUModelRunnerOutput( + model_runner_output=model_runner_output, sampled_token_ids=sampled_token_ids, + logprobs_tensors=sampler_output.logprobs_tensors, invalid_req_indices=invalid_req_indices, async_output_copy_stream=self.async_output_copy_stream, + vocab_size=self.input_batch.vocab_size, ) def _merge_additional_information_update(self, req_id: str, upd: dict) -> None: @@ -913,51 +1188,102 @@ class NPUARModelRunner(OmniNPUModelRunner): if isinstance(v, torch.Tensor): merged[k] = v.detach().to("cpu").contiguous() elif isinstance(v, list): - new_list = [] - for item in v: - if isinstance(item, torch.Tensor): - new_list.append(item.detach().to("cpu").contiguous()) - else: - new_list.append(item) - merged[k] = new_list + merged[k] = [ + (item.detach().to("cpu").contiguous() if isinstance(item, torch.Tensor) else item) for item in v + ] else: merged[k] = v setattr(req_state, "additional_information_cpu", merged) - def _generate_process_reqs_hidden_states( - self, - attn_metadata, - with_prefill, - maybe_padded_num_tokens, - input_ids, - positions, - intermediate_tensors, - inputs_embeds, - model_kwargs_extra, - ): + def _generate_process_reqs_hidden_states(self, maybe_padded_num_tokens, + input_ids, positions, + intermediate_tensors, + inputs_embeds): + model_kwargs_extra = self._build_model_kwargs_extra() + + runtime_info = model_kwargs_extra.get("runtime_additional_information", []) + if runtime_info: + for i, info in enumerate(runtime_info): + if info: + logger.debug(f"[OMNI] req[{i}] runtime_additional_information keys: {list(info.keys())}") + assert self.model is not None - hidden_states = self.model( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - **model_kwargs_extra, - ) + hidden_states = self.model(input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **self._init_model_kwargs(), + **model_kwargs_extra) forward_context = get_forward_context() - if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL: + if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL \ + and not self.use_sparse: # TODO: maybe_padded_num_tokens will be removed, use num_input_tokens instead if self.vllm_config.model_config.use_mla: - # FIXME: Try using `auto_dispatch_capture=True` - update_mla_attn_params( - self.update_stream, forward_context, maybe_padded_num_tokens, self.speculative_config - ) + if self.pcp_size * self.dcp_size > 1: + # FIXME: Try using `auto_dispatch_capture=True` + update_mla_attn_dcp_pcp_params(self.update_stream, + forward_context, + maybe_padded_num_tokens) + else: + # FIXME: Try using `auto_dispatch_capture=True` + update_mla_attn_params(self.update_stream, forward_context, + maybe_padded_num_tokens, + self.speculative_config) else: - update_attn_params(self.update_stream, forward_context, maybe_padded_num_tokens) + if self.pcp_size * self.dcp_size > 1: + update_attn_dcp_pcp_params(self.update_stream, + forward_context, + maybe_padded_num_tokens) + else: + update_attn_params(self.update_stream, forward_context, + maybe_padded_num_tokens) - if get_forward_context().sp_enabled: + if get_forward_context().sp_enabled and not isinstance( + hidden_states, IntermediateTensors): hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) pad_size = get_forward_context().pad_size if pad_size > 0: hidden_states = hidden_states[:-pad_size, :] + + if self.pcp_size > 1: + hidden_states = get_pcp_group().all_gather( + hidden_states[:self.num_actual_tokens_pcp_padded // + self.pcp_size], 0) + hidden_states = torch.index_select( + hidden_states, 0, + self.pcp_allgather_restore_idx[:hidden_states.shape[0]]) return hidden_states + + def _process_additional_information_updates( + self, + hidden_states: torch.Tensor, + multimodal_outputs: object, + num_scheduled_tokens_np: np.ndarray, + ) -> None: + """Process model-provided per-request additional_information updates and merge into request state.""" + try: + # execute the custom postprocess function + # TODO(Peiqi): do we have a more elegant way to do this? + if hasattr(self.model, "has_postprocess") and self.model.has_postprocess: + for req_index, req_id in enumerate(self.input_batch.req_ids): + req_state = self.requests.get(req_id) + req_infos = ( + getattr(req_state, "additional_information_cpu", None) if req_state is not None else None + ) + start_offset = int(self.query_start_loc.cpu[req_index]) + sched_tokens = int(num_scheduled_tokens_np[req_index]) + s, e = start_offset, start_offset + sched_tokens + # only consider to store data into update dict. + hidden_states_slice = hidden_states[s:e] + update_dict = self.model.postprocess(hidden_states_slice, **req_infos) + self._merge_additional_information_update(req_id, update_dict) + except Exception as e: + logger.error( + f"Error merging for requests:{self.input_batch.req_ids} " + f"additional information update: {e}, with the multimodal_outputs " + f"as {multimodal_outputs}" + ) + import traceback + + traceback.print_exc() diff --git a/vllm_omni/worker/npu/npu_generation_model_runner.py b/vllm_omni/worker/npu/npu_generation_model_runner.py index 33fc8e63a..836480fc8 100644 --- a/vllm_omni/worker/npu/npu_generation_model_runner.py +++ b/vllm_omni/worker/npu/npu_generation_model_runner.py @@ -3,55 +3,38 @@ from __future__ import annotations +import gc import math -from typing import TYPE_CHECKING, Any +from typing import Any import numpy as np import torch import torch.nn as nn from vllm.config import CUDAGraphMode +from vllm.distributed import get_tensor_model_parallel_world_size from vllm.distributed.parallel_state import get_pp_group -from vllm.forward_context import BatchDescriptor from vllm.logger import logger from vllm.multimodal.inputs import MultiModalKwargs from vllm.sequence import IntermediateTensors -from vllm.utils import ( - cdiv, -) +from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder from vllm.v1.core.sched.output import SchedulerOutput - -# yapf conflicts with isort for this block -# yapf: disable -from vllm.v1.kv_cache_interface import ( - EncoderOnlyAttentionSpec, -) - -# yapf: enable -from vllm.v1.outputs import ( - EMPTY_MODEL_RUNNER_OUTPUT, - AsyncModelRunnerOutput, - ModelRunnerOutput, -) +from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT from vllm.v1.spec_decode.metadata import SpecDecodeMetadata +from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorOutput -from vllm_ascend.ascend_forward_context import set_ascend_forward_context +from vllm_ascend.ascend_forward_context import MoECommType, get_mc2_tokens_capacity, set_ascend_forward_context from vllm_ascend.attention.attention_v1 import AscendAttentionState from vllm_ascend.attention.utils import AscendCommonAttentionMetadata +from vllm_ascend.patch.worker.patch_module import patch_torch_npu_argsort +from vllm_ascend.platform import NPUPlatform from vllm_ascend.spec_decode.interface import SpecDcodeType -from vllm_ascend.utils import ( - ProfileExecuteDuration, - enable_sp, - lmhead_tp_enable, -) -from vllm_ascend.worker.model_runner_v1 import AsyncNPUModelRunnerOutput +from vllm_ascend.utils import ProfileExecuteDuration, enable_sp, lmhead_tp_enable from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.worker.npu.npu_model_runner import OmniNPUModelRunner -if TYPE_CHECKING: - from vllm.v1.core.sched.output import SchedulerOutput - class NPUGenerationModelRunner(OmniNPUModelRunner): """Generation model runner for vLLM-omni on NPU (non-autoregressive).""" @@ -73,7 +56,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): torch.Tensor | None, torch.Tensor | None, int, - dict[str, Any], ]: total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens assert total_num_scheduled_tokens > 0 @@ -88,7 +70,28 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): req_ids = self.input_batch.req_ids tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids] num_scheduled_tokens = np.array(tokens, dtype=np.int32) - max_num_scheduled_tokens = num_scheduled_tokens.max() + + req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens) + _, arange = self._get_cumsum_and_arange(num_scheduled_tokens) + positions_np = np.add( + self.input_batch.num_computed_tokens_cpu[req_indices], + arange, + ) + + self.input_batch.block_table.compute_slot_mapping(req_indices, positions_np) + self.input_batch.block_table.commit_slot_mapping(total_num_scheduled_tokens) + if self.pcp_size > 1: + if not self.vllm_config.model_config.use_mla: + self.generate_kv_idx(scheduler_output) + tokens, position_pcp, pcp_unpad_mask = self._update_tokens_for_pcp(tokens) + num_scheduled_tokens = np.array(tokens, dtype=np.int32) + total_num_scheduled_tokens = sum(num_scheduled_tokens[:num_reqs]) + else: + position_pcp, pcp_unpad_mask = None, None + self.num_pcp_pads = self.num_pcp_pads[:num_reqs] + + total_num_pcp_pads = sum(self.num_pcp_pads) + max_num_scheduled_tokens = max(tokens) num_valid_tokens = np.array( [ num_tokens - len(scheduler_output.scheduled_spec_decode_tokens.get(i, [])) @@ -97,7 +100,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): dtype=np.int32, ) - if self.use_aclgraph and total_num_scheduled_tokens <= self.aclgraph_batch_sizes[-1]: + if self.use_aclgraph and total_num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]: # Add padding to the batch size. num_input_tokens = self.vllm_config.pad_for_cudagraph(total_num_scheduled_tokens) elif self.use_aclgraph and enable_sp(self.vllm_config): @@ -119,13 +122,12 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): with_prefill = attn_state not in [AscendAttentionState.DecodeOnly, AscendAttentionState.SpecDecoding] self.query_lens = torch.from_numpy(num_scheduled_tokens) - enable_dbo = self._check_dbo_is_valid(self.query_lens.tolist(), attn_state, total_num_scheduled_tokens) # Get info across DP ranks. # NOTE: maybe_padded_num_tokens is only used when using TorchAir with DP, # Otherwise, it's just max_tokens_across_dp_cpu - (maybe_padded_num_tokens, num_tokens_across_dp, with_prefill, enable_dbo) = self._sync_metadata_across_dp( - num_input_tokens, with_prefill, enable_dbo + (maybe_padded_num_tokens, num_tokens_across_dp, with_prefill) = self._sync_metadata_across_dp( + num_input_tokens, with_prefill ) # TODO: Now that num_input_tokens is basically identical with maybe_padded_num_tokens @@ -144,8 +146,15 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): # arange: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2] cu_num_tokens, arange = self._get_cumsum_and_arange(num_scheduled_tokens) - positions_np = self.positions_np[:total_num_scheduled_tokens] - np.add(self.input_batch.num_computed_tokens_cpu[req_indices], arange, out=positions_np) + if self.pcp_size > 1: + positions_np = self.positions.np[:total_num_scheduled_tokens] + np.add( + self.input_batch.num_computed_tokens_cpu[req_indices], + position_pcp[:total_num_scheduled_tokens], + out=positions_np, + ) + else: + self.positions.np[:total_num_scheduled_tokens] = positions_np # Calculate M-RoPE positions. # Only relevant for models using M-RoPE (e.g, Qwen2-VL) @@ -153,8 +162,8 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): self._calc_mrope_positions(scheduler_output) # Only relevant for models using M-RoPE (e.g, Qwen2-VL) - self.mrope_positions[:, :total_num_scheduled_tokens].copy_( - self.mrope_positions_cpu[:, :total_num_scheduled_tokens], non_blocking=True + self.mrope_positions.gpu[:, :total_num_scheduled_tokens].copy_( + self.mrope_positions.cpu[:, :total_num_scheduled_tokens], non_blocking=True ) # Get token indices. @@ -162,7 +171,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): # -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2] # where M is the max_model_len. token_indices = positions_np + req_indices * self.input_batch.token_ids_cpu.shape[1] - + token_indices_tensor = torch.from_numpy(token_indices) # Prepare input_ids. # NOTE(woosuk): We use torch.index_select instead of np.take here # because torch.index_select is much faster than np.take for large @@ -170,94 +179,184 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): torch.index_select( self.input_batch.token_ids_cpu_tensor.flatten(), 0, - torch.from_numpy(token_indices), - out=self.input_ids_cpu[:total_num_scheduled_tokens], + token_indices_tensor, + out=self.input_ids.cpu[:total_num_scheduled_tokens], ) + if self.enable_prompt_embeds: + is_token_ids = self.input_batch.is_token_ids_tensor.flatten() + torch.index_select( + is_token_ids, 0, token_indices_tensor, out=self.is_token_ids.cpu[:total_num_scheduled_tokens] + ) - # Prepare some information for building Attention-Metadata - # Compute and commit slot mapping - self.input_batch.block_table.compute_slot_mapping(req_indices, positions_np) - self.input_batch.block_table.commit_slot_mapping(total_num_scheduled_tokens) + # Because we did not pre-allocate a massive prompt_embeds CPU tensor on + # the InputBatch, we need to fill in the prompt embeds into the expected + # spots in the GpuModelRunner's pre-allocated prompt_embeds tensor. + if self.input_batch.req_prompt_embeds and (self.is_multimodal_model or self.enable_prompt_embeds): + output_idx = 0 + for req_idx in range(num_reqs): + num_sched = num_scheduled_tokens[req_idx] - self.query_start_loc_np[0] = 0 - self.query_start_loc_np[1 : num_reqs + 1] = cu_num_tokens - self.query_start_loc[: num_reqs + 1].copy_(self.query_start_loc_cpu[: num_reqs + 1], non_blocking=True) + # Skip if this request doesn't have embeddings + if req_idx not in self.input_batch.req_prompt_embeds: + output_idx += num_sched + continue - self.seq_lens_np[:num_reqs] = self.input_batch.num_computed_tokens_cpu[:num_reqs] + num_scheduled_tokens - self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs], non_blocking=True) + # Skip if no tokens scheduled + if num_sched <= 0: + output_idx += num_sched + continue + + req_embeds = self.input_batch.req_prompt_embeds[req_idx] + start_pos = self.input_batch.num_computed_tokens_cpu[req_idx] + + # Skip if trying to read beyond available embeddings + if start_pos >= req_embeds.shape[0]: + output_idx += num_sched + continue + + # Copy available embeddings + end_pos = start_pos + num_sched + actual_end = min(end_pos, req_embeds.shape[0]) + actual_num_sched = actual_end - start_pos + + if actual_num_sched > 0: + self.inputs_embeds.cpu[output_idx : output_idx + actual_num_sched].copy_( + req_embeds[start_pos:actual_end] + ) + + output_idx += num_sched + + self.query_start_loc.np[0] = 0 + self.query_start_loc.np[1 : num_reqs + 1] = cu_num_tokens + self.query_start_loc.copy_to_gpu() + + self.seq_lens.np[:num_reqs] = self.input_batch.num_computed_tokens_cpu[:num_reqs] + num_scheduled_tokens + self.seq_lens.copy_to_gpu() # Fill unused with -1. Needed for reshape_and_cache - self.query_start_loc[num_reqs + 1 :].fill_(-1) - self.seq_lens[num_reqs:].fill_(0) + self.query_start_loc.gpu[num_reqs + 1 :].fill_(-1) + self.seq_lens.gpu[num_reqs:].fill_(0) self.query_lens = torch.from_numpy(num_scheduled_tokens) # Copy the tensors to the NPU. - self._prepare_input_ids(total_num_scheduled_tokens, cu_num_tokens) - self.positions_cpu[total_num_scheduled_tokens:num_input_tokens].zero_() - self.positions[:num_input_tokens].copy_(self.positions_cpu[:num_input_tokens], non_blocking=True) + self._prepare_input_ids(scheduler_output, total_num_scheduled_tokens, cu_num_tokens) + self.positions.cpu[total_num_scheduled_tokens:num_input_tokens].zero_() + self.positions.copy_to_gpu() - # Make Attention metadata - positions_cpu = self.positions_cpu[:num_input_tokens] - positions = self.positions[:num_input_tokens] - seq_lens_cpu = self.seq_lens_cpu[:num_reqs] attn_state = self._build_attn_state(num_reqs, num_scheduled_tokens, num_valid_tokens) - self.attn_mask = self._make_attention_mask(seq_lens=seq_lens_cpu, position=positions_cpu, attn_state=attn_state) + self.attn_mask = self._make_attention_mask(attn_state) self.attn_state = attn_state # type: ignore self.with_prefill = with_prefill self.num_tokens_across_dp = num_tokens_across_dp - self._update_graph_pad_size(with_prefill, maybe_padded_num_tokens) attn_metadata: dict[str, Any] = {} + # Record the index of requests that should not be sampled, + # so that we could clear the sampled tokens before returning + num_tokens = [self.requests[r].num_tokens for r in self.input_batch.req_ids] + num_tokens_np = np.array(num_tokens, dtype=np.int32) + num_reqs = self.input_batch.num_reqs + if self.pcp_size > 1: + # while pcp > 1, we need the original num_scheduled_tokens before split + # to calculate discard_requests_mask + tokens_original = [scheduler_output.num_scheduled_tokens[i] for i in req_ids] + original_seq_lens_np = self.input_batch.num_computed_tokens_cpu[:num_reqs] + np.array( + tokens_original, dtype=np.int32 + ) + discard_requests_mask = original_seq_lens_np < num_tokens_np + else: + discard_requests_mask = self.seq_lens.np[:num_reqs] < num_tokens_np + + discard_request_indices = np.nonzero(discard_requests_mask)[0] + self.num_discarded_requests = len(discard_request_indices) + self.discard_request_indices.np[: self.num_discarded_requests] = discard_request_indices + self.discard_request_indices.copy_to_gpu(self.num_discarded_requests) + # _prepare_inputs may reorder the batch, so we must gather # multi-modal outputs after that to ensure the correct order if self.is_multimodal_model: - # Run the multimodal encoder if any. - self._execute_mm_encoder(scheduler_output) - mm_embeds = self._gather_mm_embeddings(scheduler_output) + with self.maybe_get_ec_connector_output( + scheduler_output, + encoder_cache=self.encoder_cache, + ): + # Run the multimodal encoder if any. + self._execute_mm_encoder(scheduler_output) + + # NOTE(woosuk): To unify token ids and soft tokens (vision + # embeddings), we always use embeddings (rather than token ids) + # as input to the multimodal model, even when the input is text. + input_ids = self.input_ids.gpu[:total_num_scheduled_tokens] + mm_embeds, is_mm_embed = self._gather_mm_embeddings(scheduler_output) + + inputs_embeds = self.model.embed_input_ids( + input_ids, + multimodal_embeddings=mm_embeds, + is_multimodal=is_mm_embed, + ) - # -------------------------------------- Omni-new ------------------------------------------------- - # NOTE(woosuk): To unify token ids and soft tokens (vision - # embeddings), we always use embeddings (rather than token ids) - # as input to the multimodal model, even when the input is text. - input_ids = self.input_ids[:num_input_tokens] - if mm_embeds: - inputs_embeds = self.model.get_input_embeddings(input_ids, mm_embeds) - else: - inputs_embeds = self.model.get_input_embeddings(input_ids) # TODO(woosuk): Avoid the copy. Optimize. - self.inputs_embeds[:num_input_tokens].copy_(inputs_embeds) - inputs_embeds = self.inputs_embeds[:num_input_tokens] - + self.inputs_embeds.gpu[:total_num_scheduled_tokens].copy_(inputs_embeds) + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] + # -------------------------------------- Omni-new ------------------------------------------------- + # NOTE(gcanlin): We don't set input_ids to None in vllm-omni. model_kwargs = { - **self._init_model_kwargs(num_input_tokens), + **self._init_model_kwargs(), **self._extract_mm_kwargs(scheduler_output), } - # (NOTE) Omni-new: input_ids isn't set as None - # ------------------------------------------------------------------------------------------------- + # -------------------------------------- Omni-new ------------------------------------------------- + elif self.enable_prompt_embeds and get_pp_group().is_first_rank: + # Get the input embeddings for the tokens that are not input embeds, + # then put them into the appropriate positions. + # TODO(qthequartermasterman): Since even when prompt embeds are + # enabled, (a) not all requests will use prompt embeds, and (b) + # after the initial prompt is processed, the rest of the generated + # tokens will be token ids, it is not desirable to have the + # embedding layer outside of the acl graph all the time. The v0 + # engine avoids this by "double compiling" the acl graph, once + # with input_ids and again with inputs_embeds, for all num_tokens. + # If a batch only has token ids, then including the embedding layer + # in the acl graph will be more performant (like in the else case + # below). + token_ids_idx = self.is_token_ids.gpu[:total_num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1) + # Some tokens ids may need to become embeds + if token_ids_idx.numel() > 0: + token_ids = self.input_ids.gpu[token_ids_idx] + tokens_to_embeds = self.model.embed_input_ids(input_ids=token_ids) + self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds + + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] + input_ids = None else: # For text-only models, we use token ids as input. # While it is possible to use embeddings as input just like the # multimodal models, it is not desirable for performance since # then the embedding layer is not included in the ACL graph. - input_ids = self.input_ids[:num_input_tokens] + input_ids = self.input_ids.gpu[:num_input_tokens] inputs_embeds = None - model_kwargs = self._init_model_kwargs(num_input_tokens) - positions = self.positions[:num_input_tokens] - input_ids, positions = self._update_input_ids_and_positions( - input_ids, positions, num_input_tokens, with_prefill, maybe_padded_num_tokens - ) + positions = self.positions.gpu[:num_input_tokens] + if self.uses_mrope: + positions = self.mrope_positions.gpu[:, :num_input_tokens] + # type: ignore if get_pp_group().is_first_rank: intermediate_tensors = None else: assert intermediate_tensors is not None assert self.intermediate_tensors is not None + # If both flashcomm1 and pp are used simultaneously, + # the shape of the received data and the shape of the space to be copied to will not match, + # requiring a recalculation of the incoming data's shape. + tp_size = get_tensor_model_parallel_world_size() + num_input_tokens_with_flashcomm1 = num_input_tokens + if enable_sp(): + num_input_tokens_with_flashcomm1 = (num_input_tokens + tp_size - 1) // tp_size for k, v in intermediate_tensors.items(): - self.intermediate_tensors[k][:num_input_tokens].copy_(v[:num_input_tokens], non_blocking=True) + self.intermediate_tensors[k][:num_input_tokens_with_flashcomm1].copy_( + v[:num_input_tokens_with_flashcomm1], non_blocking=True + ) intermediate_tensors = IntermediateTensors( - {k: v[:num_input_tokens] for k, v in self.intermediate_tensors.items()} + {k: v[:num_input_tokens_with_flashcomm1] for k, v in self.intermediate_tensors.items()} ) use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 @@ -268,34 +367,66 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): # We will ignore the sampled tokens from the partial requests. # TODO: Support prompt logprobs. spec_decode_metadata = None - logits_indices = torch.from_numpy(cu_num_tokens - 1).to(self.device, non_blocking=True) + if self.pcp_size * self.dcp_size > 1: + logits_indices = torch.from_numpy(cu_num_tokens) * self.pcp_size - self.num_pcp_pads[:num_reqs] - 1 + logits_indices = logits_indices.pin_memory().to(self.device, non_blocking=True) + else: + logits_indices = self.query_start_loc.gpu[1 : num_reqs + 1] - 1 else: # Get the number of draft tokens for each request. # Iterate over the dictionary rather than all requests since not all # requests have draft tokens. num_draft_tokens = np.zeros(num_reqs, dtype=np.int32) + # For chunked prefills, use -1 as mask rather than 0, as guided + # decoding may rollback speculative tokens. + num_decode_draft_tokens = np.full(num_reqs, -1, dtype=np.int32) for req_id, draft_token_ids in scheduler_output.scheduled_spec_decode_tokens.items(): req_idx = self.input_batch.req_id_to_index[req_id] num_draft_tokens[req_idx] = len(draft_token_ids) + num_decode_draft_tokens[req_idx] = ( + len(draft_token_ids) + if ( + self.input_batch.num_computed_tokens_cpu[req_idx] >= self.input_batch.num_prompt_tokens[req_idx] + ) + else -1 + ) - spec_decode_metadata = self._calc_spec_decode_metadata(num_draft_tokens, cu_num_tokens) + spec_decode_metadata = self._calc_spec_decode_metadata( + num_draft_tokens, cu_num_tokens, self.num_pcp_pads[:num_reqs] + ) logits_indices = spec_decode_metadata.logits_indices - self.num_draft_tokens.np[:num_reqs] = num_draft_tokens - self.num_draft_tokens.np[num_reqs:].fill(0) - self.num_draft_tokens.copy_to_gpu() + + # For DECODE only cuda graph of some attention backends (e.g., GDN). + self.num_decode_draft_tokens.np[:num_reqs] = num_decode_draft_tokens + self.num_decode_draft_tokens.np[num_reqs:].fill(-1) + self.num_decode_draft_tokens.copy_to_gpu() + # save logits_indices for pcp spec decode usage + self.logits_indices = logits_indices # Used in the below loop. # query_start_loc_cpu = self.query_start_loc.cpu[:num_reqs + 1] num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs] - spec_decode_common_attn_metadata = None + self.spec_decode_common_attn_metadata = None if use_spec_decode and self.need_accepted_tokens: self.num_accepted_tokens.np[:num_reqs] = self.input_batch.num_accepted_tokens_cpu[:num_reqs] self.num_accepted_tokens.np[num_reqs:].fill(1) self.num_accepted_tokens.copy_to_gpu() + if self.speculative_config and self.pcp_size > 1: + self._generate_pcp_mtp_input( + num_reqs, scheduler_output.total_num_scheduled_tokens, scheduler_output.num_scheduled_tokens + ) + + long_seq_metadata = self._generate_pcp_metadata(total_num_scheduled_tokens) # Prepare the attention metadata for each KV cache group and make layers # in the same group share the same metadata. for kv_cache_group_id, kv_cache_group_spec in enumerate(self.kv_cache_config.kv_cache_groups): + # NOTE: This is strange, why did we use total_num_scheduled_tokens before? + slot_mapping_size = ( + total_num_scheduled_tokens + if self.pcp_size == 1 + else total_num_scheduled_tokens * self.pcp_size - total_num_pcp_pads + ) if isinstance(kv_cache_group_spec.kv_cache_spec, EncoderOnlyAttentionSpec): # Encoder-only layers do not have KV cache, so we need to # create a dummy block table and slot mapping for them. @@ -312,58 +443,116 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): else: blk_table = self.input_batch.block_table[kv_cache_group_id] blk_table_tensor = blk_table.get_device_tensor() - slot_mapping = blk_table.slot_mapping_cpu[:total_num_scheduled_tokens] - self.slot_mapping[:total_num_scheduled_tokens].copy_( - slot_mapping[:total_num_scheduled_tokens], - non_blocking=True, - ) - self.slot_mapping[total_num_scheduled_tokens:].fill_(0) + blk_table.slot_mapping.gpu[slot_mapping_size:].fill_(0) + if self.pcp_size > 1: + slot_mapping_for_pcp = blk_table.slot_mapping.gpu[: long_seq_metadata.num_actual_tokens_pcp_padded] + slot_mapping_for_pcp[slot_mapping_size:].fill_(-1) + assert pcp_unpad_mask is not None + pcp_padded_slot_mapping = self.pcp_padded_slot_mapping[: pcp_unpad_mask.shape[0]] + pcp_padded_slot_mapping.fill_(-1) + pcp_padded_slot_mapping[pcp_unpad_mask] = slot_mapping_for_pcp[:slot_mapping_size] + slot_mapping_for_pcp[: long_seq_metadata.num_actual_tokens_pcp_padded] = pcp_padded_slot_mapping + blk_table.slot_mapping.gpu[: long_seq_metadata.num_actual_tokens_pcp_padded] = slot_mapping_for_pcp + slot_mapping = blk_table.slot_mapping.gpu + + # NOTE: This is a temporary hack, now in GPUModelRunner, this prepare_inputs + # has been split to multiple parts, and there are 3 parts that is related to this + # `num_reqs`, we'll take `query_start_loc` as an example: + # 1. self.query_start_loc.np[1 : num_reqs + 1] = cu_num_tokens + # 2. get `num_reqs_padded`, this depends on dispatcher and which is why we have the + # following simplified `dispatch` logic here, we try to minimize the impact + # 3. query_start_loc = self.query_start_loc.gpu[: num_reqs_padded + 1] + uniform_decode = (max_num_scheduled_tokens == self.uniform_decode_query_len) and ( + total_num_scheduled_tokens == max_num_scheduled_tokens * num_reqs + ) + + # TODO: We should make this official ASAP. Also note that if we pad here, + # the builders won’t need to add any extra padding. + if self.compilation_config.cudagraph_mode.decode_mode() == CUDAGraphMode.FULL and uniform_decode: + num_reqs_padded = num_input_tokens // self.uniform_decode_query_len + pad_size = num_reqs_padded - num_reqs + if pad_size > 0: + last_query_loc = self.query_start_loc.gpu[num_reqs] + + steps = torch.arange(1, pad_size + 1, device=self.device, dtype=self.query_start_loc.gpu.dtype) + fill_values = last_query_loc + (steps * self.uniform_decode_query_len) + + self.query_start_loc.gpu[num_reqs + 1 : num_reqs_padded + 1] = fill_values + # So we are trying to simulate the behavior of GPUModelRunner's + # prepare_inputs for uniform decode mode by padding query_start_loc + num_reqs = num_reqs_padded # Make AscendCommonAttentionMetadata common_attn_metadata = AscendCommonAttentionMetadata( - query_start_loc=self.query_start_loc[: num_reqs + 1], - query_start_loc_cpu=self.query_start_loc_cpu[: num_reqs + 1], - seq_lens_cpu=self.seq_lens_cpu, - seq_lens=self.seq_lens_cpu[:num_reqs], + query_start_loc=self.query_start_loc.gpu[: num_reqs + 1], + query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1], + seq_lens_cpu=self.seq_lens.cpu[:num_reqs], + seq_lens=self.seq_lens.gpu[:num_reqs], num_reqs=num_reqs, - num_actual_tokens=total_num_scheduled_tokens, + num_actual_tokens=slot_mapping_size, num_input_tokens=num_input_tokens, actual_seq_lengths_q=self.actual_seq_lengths_q, # TODO: change this to the right block table for linear attn block_table_tensor=blk_table_tensor[:num_reqs], - slot_mapping=self.slot_mapping, + slot_mapping=slot_mapping, num_computed_tokens_cpu=num_computed_tokens_cpu, - positions=self.positions, + positions=self.positions.gpu, attn_mask=self.attn_mask, spec_attn_mask=self.spec_attn_mask, attn_state=self.attn_state, - enable_dbo_across_dp=enable_dbo, is_only_prefill=bool(np.all(num_valid_tokens != 1)), max_query_len=max_num_scheduled_tokens, - graph_pad_size=self.graph_pad_size, decode_token_per_req=self.decode_token_per_req, - cos=self.cos, - sin=self.sin, + prefill_context_parallel_metadata=long_seq_metadata, ) - if self.speculative_config and spec_decode_common_attn_metadata is None: - spec_decode_common_attn_metadata = common_attn_metadata + if self.speculative_config and self.pcp_size > 1: + # For pcp + spec decode, we flatten block_table + # to avoid irregular spec_attn_mask shape, e.g., + # num_decode_req=2, num_prefill_req=3, num_speculative_tokens=1, + # ori block_table: # [d0, d1, p0, p1, p2] + # (num_reqs_d + num_reqs_p, max_num_blocks), + # flattened block_table: [d0, d0, d1, d1, p0, p1, p2] + # (num_reqs_d * decode_threshold + num_reqs_p, max_num_blocks), + ori_query_lens = ( + self.query_start_loc_pcp_full.cpu[1 : num_reqs + 1] - self.query_start_loc_pcp_full.cpu[:num_reqs] + ) + num_prefill_reqs = (ori_query_lens > self.decode_threshold).sum().item() + num_decode_reqs = num_reqs - num_prefill_reqs + num_decode_reqs_flatten = num_decode_reqs * self.decode_threshold + blk_table_tensor[num_decode_reqs_flatten : num_decode_reqs_flatten + num_prefill_reqs].copy_( + blk_table_tensor[num_decode_reqs : num_decode_reqs + num_prefill_reqs].clone() + ) + blk_table_tensor[:num_decode_reqs_flatten].copy_( + blk_table_tensor[:num_decode_reqs].repeat_interleave(self.decode_threshold, dim=0) + ) + common_attn_metadata.block_table_tensor = blk_table_tensor[: num_decode_reqs_flatten + num_prefill_reqs] + + if self.speculative_config and self.spec_decode_common_attn_metadata is None: + self.spec_decode_common_attn_metadata = common_attn_metadata for attn_group in self.attn_groups[kv_cache_group_id]: common_prefix_len = 0 extra_attn_metadata_args = {} builder = attn_group.get_metadata_builder() - if isinstance(builder, GDNAttentionMetadataBuilder) or self.model_config.runner_type == "pooling": + if isinstance(builder, GDNAttentionMetadataBuilder): if use_spec_decode: + patch_torch_npu_argsort() extra_attn_metadata_args = dict( num_accepted_tokens=self.num_accepted_tokens.gpu[:num_reqs], - num_draft_tokens=self.num_draft_tokens.gpu[:num_reqs], + num_decode_draft_tokens_cpu=self.num_decode_draft_tokens.cpu[:num_reqs], ) attn_metadata_i = builder.build( common_prefix_len=common_prefix_len, common_attn_metadata=common_attn_metadata, **extra_attn_metadata_args, ) + elif self.model_config.runner_type == "pooling": + attn_metadata_i = builder.build( + common_prefix_len=common_prefix_len, + common_attn_metadata=common_attn_metadata, + **extra_attn_metadata_args, + ) else: attn_metadata_i = builder.build( common_prefix_len=common_prefix_len, @@ -376,13 +565,13 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): attn_metadata[layer_name] = attn_metadata_i if lmhead_tp_enable(): - max_num_reqs_across_dp = maybe_padded_num_tokens if not with_prefill else self.max_num_reqs + max_num_reqs_across_dp = self.max_num_reqs * self.uniform_decode_query_len logits_indices = nn.functional.pad(logits_indices, (0, max_num_reqs_across_dp - logits_indices.shape[0])) return ( attn_metadata, positions, - num_input_tokens, + num_scheduled_tokens, num_input_tokens, num_tokens_across_dp, maybe_padded_num_tokens, @@ -400,7 +589,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): self, scheduler_output: SchedulerOutput, intermediate_tensors: IntermediateTensors | None = None, - ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: + ) -> OmniModelRunnerOutput | IntermediateTensors: with ProfileExecuteDuration().capture_async("prepare input"): self._update_states(scheduler_output) if not scheduler_output.total_num_scheduled_tokens: @@ -428,13 +617,25 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): if self.dynamic_eplb: self.eplb_updator.take_update_info_from_eplb_process() - moe_comm_type = self._select_moe_comm_method(num_input_tokens, self.with_prefill) + moe_comm_type = self._select_moe_comm_method(num_input_tokens) + # prevent debugger is None + need_dump = self.dump_enable and self.debugger is not None + if need_dump: + assert self.debugger is not None + dbg_cfg = getattr(self.debugger, "config", None) + dump_level = str(getattr(dbg_cfg, "level", "L1")).upper() if dbg_cfg is not None else "L1" + if dump_level in ("L0", "MIX"): + self.debugger.start(model=self.model) + else: + self.debugger.start() - # -------------------------------------- Omni-new ------------------------------------------------- - # Omni-new: don't use cudagraph_dispatcher - # and remove ubatch_slices - aclgraph_runtime_mode = CUDAGraphMode.NONE - # ------------------------------------------------------------------------------------------------- + uniform_decode = (max_query_len == self.uniform_decode_query_len) and ( + scheduler_output.total_num_scheduled_tokens == self.input_batch.num_reqs * max_query_len + ) + has_lora = len(self.input_batch.lora_id_to_lora_request) > 0 + aclgraph_runtime_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch( + num_tokens=num_input_tokens, uniform_decode=uniform_decode, has_lora=has_lora + ) # Run forward pass with ProfileExecuteDuration().capture_async("forward"): @@ -444,17 +645,16 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): num_tokens=num_input_tokens, num_tokens_across_dp=num_tokens_across_dp, with_prefill=self.with_prefill, - reserved_mc2_mask=self.reserved_mc2_mask, moe_comm_type=moe_comm_type, aclgraph_runtime_mode=aclgraph_runtime_mode, - batch_descriptor=None, + batch_descriptor=batch_descriptor, num_actual_tokens=scheduler_output.total_num_scheduled_tokens, prefetch_stream=self.prefetch_stream, model_instance=self.model, weight_prefetch_method=self.weight_prefetch_method, ): self.maybe_setup_kv_connector(scheduler_output) - + # -------------------------------------- Omni-new ------------------------------------------------- outputs = self._run_generation( input_ids=input_ids, positions=positions, @@ -463,6 +663,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): multimodal_kwargs=model_kwargs, logits_indices=logits_indices, ) + # -------------------------------------- Omni-new ------------------------------------------------- self.maybe_wait_for_kv_save() finished_sending, finished_recving = self.get_finished_kv_transfer(scheduler_output) @@ -474,26 +675,30 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): kv_connector_output = KVConnectorOutput(finished_sending=finished_sending, finished_recving=finished_recving) finished_sending = None finished_recving = None - - # -------------------------------------- Omni-new ------------------------------------------------- - # Omni-new: extract_multimodal_outputs + # -------------------------------------- Omni-new ------------------------------------------------- + # We don't need any post-process for generation model outputs _, multimodal_outputs = self.extract_multimodal_outputs(outputs) - # Ensure one tensor per request, map to CPU for output struct - pooler_output: list[torch.Tensor | None] = [] + pooler_output: list[object] = [] if isinstance(multimodal_outputs, torch.Tensor): - # If model returned a single stacked tensor, split by requests + assert multimodal_outputs.shape[0] == 1, ( + "model should return a single tensor, to return multiple tensors, use a dict" + ) assert multimodal_outputs.shape[0] == self.input_batch.num_reqs for i in range(self.input_batch.num_reqs): - pooler_output.append(multimodal_outputs[i].detach().to("cpu").contiguous()) + pooler_output.append({"model_outputs": multimodal_outputs[i].detach().to("cpu").contiguous()}) elif isinstance(multimodal_outputs, list): + assert len(multimodal_outputs) == 1, ( + "model should return a single list, to return multiple lists, use a dict" + ) for out in multimodal_outputs: - pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None) + pooler_output.append( + {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None} + ) elif isinstance(multimodal_outputs, dict): - for out in multimodal_outputs.values(): - pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None) + for key, out in multimodal_outputs.items(): + pooler_output.append({key: out.detach().to("cpu").contiguous() if out is not None else None}) else: raise RuntimeError("Unsupported diffusion output type") - output = OmniModelRunnerOutput( req_ids=self.input_batch.req_ids, req_id_to_index=self.input_batch.req_id_to_index, @@ -504,21 +709,11 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): kv_connector_output=kv_connector_output, num_nans_in_logits={}, ) - # ------------------------------------------------------------------------------------------------- - - durations = ProfileExecuteDuration().pop_captured_sync() - if durations: - dr_str = [f"[{tag}]:{duration:.2f}ms" for tag, duration in durations.items()] - captured_name = "Decode" if self.attn_state == AscendAttentionState.DecodeOnly else "Prefill" - logger.info("Profile execute duration [%s]:%s", captured_name, " ".join(dr_str)) - - if self.dynamic_eplb: - self.eplb_updator.forward_end() - + # -------------------------------------- Omni-new ------------------------------------------------- if not self.use_async_scheduling: return output - return AsyncNPUModelRunnerOutput( + return AsyncGPUModelRunnerOutput( model_runner_output=output, sampled_token_ids=[], invalid_req_indices=[], @@ -535,12 +730,14 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): multimodal_kwargs: dict, logits_indices: torch.Tensor, ) -> torch.Tensor | list[torch.Tensor]: - """Runs the generation process and returns per-request tensors. + """Run generation from codec codes to waveforms. - Tries model interfaces in the following order for maximal compatibility: - 1) model.sample(condition=..., **kwargs) - 2) model.forward(condition=..., **kwargs) - 3) model.diffuse(condition=..., **kwargs) + Args: + scheduler_output: Contains codec codes in input_ids or additional info + intermediate_tensors: PP intermediate tensors if applicable + + Returns: + Audio waveforms: [batch, 1, waveform_len] or list of tensors """ kwargs = dict( input_ids=input_ids, @@ -555,19 +752,22 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): if hasattr(self.model, "forward"): return self.model.forward(**kwargs) - # TODO: add the diffuse method for other models raise RuntimeError( "The loaded model does not expose generation interfaces 'sample', " "'forward', or 'diffuse'. Please implement one of them or adapt the runner." ) + @torch.inference_mode() + def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None: + logger.warning("Dummy sampler run is not implemented for generation model") + return None + @torch.inference_mode() def _dummy_run( self, num_tokens: int, with_prefill: bool = False, - is_torchair_compile: bool = False, aclgraph_runtime_mode: CUDAGraphMode | None = None, force_attention: bool = False, uniform_decode: bool = False, @@ -578,7 +778,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL, } - # In multi-DP scenarios, there may be situations where all DP groups are executing dummy runs. # If sequence parallelism is enabled, it is essential to ensure that num_tokens is divisible by tp_size. if self.use_aclgraph and enable_sp(self.vllm_config): @@ -590,11 +789,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): with_prefill = True # Padding for DP - (num_tokens, num_tokens_across_dp, with_prefill, _) = self._sync_metadata_across_dp( - num_tokens, with_prefill, False - ) - - moe_comm_type = self._select_moe_comm_method(num_tokens, with_prefill) + (num_tokens, num_tokens_across_dp, with_prefill) = self._sync_metadata_across_dp(num_tokens, with_prefill) # If cudagraph_mode.decode_mode() == FULL and # cudagraph_mode.separate_routine(). This means that we are using @@ -633,76 +828,109 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): assert sum(num_scheduled_tokens_list) == num_tokens assert len(num_scheduled_tokens_list) == num_reqs num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32) + num_sampled_tokens = np.ones(num_reqs, dtype=np.int32) if not self.in_profile_run and self.dynamic_eplb: self.eplb_updator.forward_before() - with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens): + has_lora = True if self.lora_config and self.compilation_config.cudagraph_specialize_lora else False + _ag_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch( + num_tokens=num_tokens, uniform_decode=uniform_decode, has_lora=has_lora + ) + + num_tokens_padded = batch_descriptor.num_tokens + num_reqs_padded = batch_descriptor.num_reqs if batch_descriptor.num_reqs is not None else num_reqs + if num_tokens_across_dp is not None and num_tokens_padded != num_tokens: + # pad is needed if the pad of `num_tokens` is triggered inside CudagraphDispatcher + num_tokens_across_dp[:] = num_tokens_padded + num_scheduled_tokens = num_scheduled_tokens.repeat(num_reqs_padded) + + moe_comm_type = self._select_moe_comm_method(num_tokens_padded) + + # filter out the valid batch descriptor + if aclgraph_runtime_mode is not None: + # we allow forcing NONE when the dispatcher disagrees to support + # warm ups for aclgraph capture + if aclgraph_runtime_mode != CUDAGraphMode.NONE and aclgraph_runtime_mode != _ag_mode: + raise ValueError( + f"Aclgraph runtime mode mismatch at dummy_run. " + f"Expected {_ag_mode}, but got {aclgraph_runtime_mode}." + ) + else: + aclgraph_runtime_mode = _ag_mode + + # TODO(Mengqing): Set create_mixed_batch to False since it's only used in FI warmup + # and not supported in ASCEND now. We could remove it in the future. + attn_metadata = self._build_dummy_attn_metadata( + False, + num_reqs=num_reqs_padded, + num_tokens=num_tokens_padded, + max_query_len=max_query_len, + aclgraph_runtime_mode=aclgraph_runtime_mode, + force_attention=force_attention, + num_scheduled_tokens=num_scheduled_tokens, + ) + + with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, num_sampled_tokens): + # Make sure padding doesn't exceed max_num_tokens + assert num_tokens_padded <= self.max_num_tokens if self.is_multimodal_model: input_ids = None - inputs_embeds = self.inputs_embeds[:num_tokens] + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] + elif self.enable_prompt_embeds: + input_ids = None + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] else: - input_ids = self.input_ids[:num_tokens] + input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None if self.uses_mrope: - positions = self.mrope_positions[:, :num_tokens] + positions = self.mrope_positions.gpu[:, :num_tokens_padded] else: - positions = self.positions[:num_tokens] + positions = self.positions.gpu[:num_tokens_padded] if get_pp_group().is_first_rank: intermediate_tensors = None else: + # When PP and flashcomm1 are enabled, + # during dummy_run the estimated space should divide num_tokens by tp_size; + # otherwise, on non-first PP ranks it would effectively perform an extra all-gather, + # leading to incorrect memory estimation and potentially causing OOM. + actual_tokens = num_tokens + if enable_sp(): + tp_size = get_tensor_model_parallel_world_size() + actual_tokens = num_tokens // tp_size if self.intermediate_tensors is None: self.intermediate_tensors = self.model.make_empty_intermediate_tensors( - batch_size=num_tokens, dtype=self.dtype, device=self.device + batch_size=actual_tokens, dtype=self.dtype, device=self.device ) intermediate_tensors = IntermediateTensors( - {k: v[:num_tokens] for k, v in self.intermediate_tensors.items()} + {k: v[:num_tokens_padded] for k, v in self.intermediate_tensors.items()} ) - # filter out the valid batch descriptor - _ag_mode, batch_descriptor = self.aclgraph_dispatcher.dispatch( - BatchDescriptor(num_tokens=num_tokens, uniform_decode=uniform_decode) - ) - if aclgraph_runtime_mode is not None: - # we allow forcing NONE when the dispatcher disagrees to support - # warm ups for aclgraph capture - assert aclgraph_runtime_mode == CUDAGraphMode.NONE or aclgraph_runtime_mode == _ag_mode, ( - f"Aclgraph runtime mode mismatch at dummy_run. " - f"Expected {_ag_mode}, but got {aclgraph_runtime_mode}." - ) - else: - aclgraph_runtime_mode = _ag_mode - - # TODO(Mengqing): Set create_mixed_batch to False since it's only used in FI warmup - # and not supported in ASCEND now. We could remove it in the future. - attn_metadata = self._build_dummy_attn_metadata( - False, - num_reqs=num_reqs, - num_tokens=num_tokens, - max_query_len=max_query_len, - aclgraph_runtime_mode=aclgraph_runtime_mode, - force_attention=force_attention, - ) - need_dummy_logits = not self.in_profile_run and lmhead_tp_enable() + max_num_reqs_across_dp = max_num_reqs * self.uniform_decode_query_len + dummy_indices = torch.zeros(max_num_reqs_across_dp, dtype=torch.int32) - if need_dummy_logits: - max_num_reqs_across_dp = num_tokens if not with_prefill else max_num_reqs - dummy_indices = torch.zeros(max_num_reqs_across_dp, dtype=torch.int32) + def dummy_compute_logits(hidden_states): + if not need_dummy_logits: + return None + return self.model.compute_logits(hidden_states[dummy_indices]) - def dummy_compute_logits(hidden_states): - return self.model.compute_logits(hidden_states[dummy_indices]) + def dummy_drafter_compute_logits(hidden_states): + if not need_dummy_logits or self.drafter is None: + return + if hasattr(self.drafter, "model") and hasattr(self.drafter.model, "compute_logits"): + return self.drafter.model.compute_logits(hidden_states[dummy_indices]) with set_ascend_forward_context( attn_metadata, self.vllm_config, - num_tokens=num_tokens, + num_tokens=num_tokens_padded, num_tokens_across_dp=num_tokens_across_dp, with_prefill=with_prefill, in_profile_run=self.in_profile_run, - reserved_mc2_mask=self.reserved_mc2_mask, + # reserved_mc2_mask=self.reserved_mc2_mask, moe_comm_type=moe_comm_type, num_actual_tokens=0, aclgraph_runtime_mode=aclgraph_runtime_mode, @@ -712,37 +940,50 @@ class NPUGenerationModelRunner(OmniNPUModelRunner): weight_prefetch_method=self.weight_prefetch_method, ): hidden_states = self._generate_dummy_run_hidden_states( - with_prefill, - is_torchair_compile, - input_ids, - positions, - attn_metadata, - num_tokens, - intermediate_tensors, - inputs_embeds, + input_ids, positions, num_tokens_padded, intermediate_tensors, inputs_embeds ) - if need_dummy_logits: - dummy_compute_logits(hidden_states) + dummy_compute_logits(hidden_states) if self.drafter: self.drafter.dummy_run( - num_tokens=num_tokens, + num_tokens=num_tokens_padded, with_prefill=with_prefill, - skip_attn=True, - num_reqs=num_reqs, + num_reqs=num_reqs_padded, num_tokens_across_dp=num_tokens_across_dp, aclgraph_runtime_mode=aclgraph_runtime_mode, batch_descriptor=batch_descriptor, + dummy_compute_logits=dummy_drafter_compute_logits, + skip_attn=not force_attention, ) - if need_dummy_logits: - self.drafter.model.compute_logits(hidden_states[dummy_indices]) if self.in_profile_run and self.dynamic_eplb: self.model.clear_all_moe_loads() if not self.in_profile_run and self.dynamic_eplb: self.eplb_updator.take_update_info_from_eplb_process() self.eplb_updator.forward_end() - # -------------------------------------- Omni-new ------------------------------------------------- hidden_states, _ = self.extract_multimodal_outputs(hidden_states) # ------------------------------------------------------------------------------------------------- return hidden_states + + def profile_run(self) -> None: + # Trigger compilation for general shape. + with self.set_in_profile_run(): + hidden_states = self._dummy_run( + self.max_num_tokens // self.pcp_size if self.pcp_size > 1 else self.max_num_tokens, with_prefill=True + ) + # MC2 will consume additional NPU memory. + # Therefore, we need to run the MC2 path once here to complete its initialization, + # allowing vLLM to correctly estimate the maximum memory required. + mc2_tokens_capacity = get_mc2_tokens_capacity() + if ( + self.max_num_tokens > mc2_tokens_capacity + and self._select_moe_comm_method(mc2_tokens_capacity) == MoECommType.MC2 + ): + self._dummy_run(mc2_tokens_capacity, with_prefill=True) + + output = None + + NPUPlatform.synchronize() + del hidden_states, output + self.encoder_cache.clear() + gc.collect() diff --git a/vllm_omni/worker/npu/npu_model_runner.py b/vllm_omni/worker/npu/npu_model_runner.py index c57c6df48..e1fb9a7a7 100644 --- a/vllm_omni/worker/npu/npu_model_runner.py +++ b/vllm_omni/worker/npu/npu_model_runner.py @@ -2,27 +2,24 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, cast import numpy as np import torch from vllm.config import CUDAGraphMode +from vllm.distributed import get_tensor_model_parallel_world_size from vllm.distributed.parallel_state import get_pp_group -from vllm.forward_context import BatchDescriptor from vllm.logger import init_logger from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding -from vllm.model_executor.models.interfaces import SupportsMultiModal, supports_mrope +from vllm.model_executor.models.interfaces import supports_mrope from vllm.model_executor.models.interfaces_base import VllmModelForPooling -from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.inputs import BatchedTensorInputs, MultiModalKwargsItem -from vllm.multimodal.utils import group_mm_kwargs_by_modality from vllm.sampling_params import SamplingType -from vllm.utils import cdiv -from vllm.v1.worker.gpu_model_runner import IntermediateTensors +from vllm.sequence import IntermediateTensors +from vllm.utils.math_utils import cdiv +from vllm.v1.worker.gpu_input_batch import CachedRequestState from vllm_ascend.ascend_forward_context import set_ascend_forward_context from vllm_ascend.utils import enable_sp, lmhead_tp_enable from vllm_ascend.worker.model_runner_v1 import NPUModelRunner -from vllm_ascend.worker.npu_input_batch import CachedRequestState if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput @@ -37,10 +34,9 @@ class OmniNPUModelRunner(NPUModelRunner): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - - self.is_multimodal_raw_input_only_model = self.model_config.is_multimodal_raw_input_only_model - self.mm_registry = MULTIMODAL_REGISTRY - self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(self.model_config) + self._omni_per_req_additional_information: dict[str, dict] | None = None + self._omni_num_scheduled_tokens_np: np.ndarray | None = None + self._omni_last_model_output: object | None = None def _init_mrope_positions(self, req_state: CachedRequestState): image_grid_thw = [] @@ -88,10 +84,18 @@ class OmniNPUModelRunner(NPUModelRunner): ) def _update_states(self, scheduler_output: "SchedulerOutput") -> None: + """Update the cached states and the persistent batch with the scheduler + output. + + The updated states are used by the `_prepare_inputs` function to create + the input GPU tensors for the model. + + The SamplingMetadata is updated and copied to the GPU if there is a + new/resumed/paused/finished request in the batch. + """ # Remove finished requests from the cached states. for req_id in scheduler_output.finished_req_ids: self.requests.pop(req_id, None) - # Remove the finished requests from the persistent batch. # NOTE(woosuk): There could be an edge case where finished_req_ids and # scheduled_req_ids overlap. This happens when a request is aborted and @@ -100,8 +104,11 @@ class OmniNPUModelRunner(NPUModelRunner): # and handling the second as a new request. for req_id in scheduler_output.finished_req_ids: self.input_batch.remove_request(req_id) + + # Free the cached encoder outputs. for mm_hash in scheduler_output.free_encoder_mm_hashes: self.encoder_cache.pop(mm_hash, None) + # Remove the unscheduled requests from the persistent batch. # NOTE(woosuk): The unscheduled requests are either preempted requests # or running requests that are not scheduled in this step. We remove @@ -117,7 +124,7 @@ class OmniNPUModelRunner(NPUModelRunner): for req_id in unscheduled_req_ids: self.input_batch.remove_request(req_id) - req_ids_to_add: list[str] = [] + reqs_to_add: list[CachedRequestState] = [] # Add new requests to the cached states. for new_req_data in scheduler_output.scheduled_new_reqs: req_id = new_req_data.req_id @@ -130,18 +137,20 @@ class OmniNPUModelRunner(NPUModelRunner): else: generator = None - if pooling_params: - assert (task := pooling_params.task) is not None, "You did not set `task` in the API" + if self.is_pooling_model: + assert pooling_params is not None + task = pooling_params.task + assert task is not None, "You did not set `task` in the API" + model = cast(VllmModelForPooling, self.get_model()) to_update = model.pooler.get_pooling_updates(task) to_update.apply(pooling_params) - backward_kwargs = {} - backward_kwargs["mm_features"] = new_req_data.mm_features - - self.requests[req_id] = CachedRequestState( + req_state = CachedRequestState( req_id=req_id, prompt_token_ids=new_req_data.prompt_token_ids, + prompt_embeds=new_req_data.prompt_embeds, + mm_features=new_req_data.mm_features, sampling_params=sampling_params, pooling_params=pooling_params, generator=generator, @@ -149,8 +158,8 @@ class OmniNPUModelRunner(NPUModelRunner): num_computed_tokens=new_req_data.num_computed_tokens, output_token_ids=[], lora_request=new_req_data.lora_request, - **backward_kwargs, ) + self.requests[req_id] = req_state # -------------------------------------- Omni-new ------------------------------------------------- # If prompt embeddings are provided, decode and attach to inter_data @@ -187,7 +196,7 @@ class OmniNPUModelRunner(NPUModelRunner): dt = np.dtype(getattr(entry, "tensor_dtype", "float32")) arr = np.frombuffer(entry.tensor_data, dtype=dt) arr = arr.reshape(entry.tensor_shape) - info_dict[k] = torch.from_numpy(arr) + info_dict[k] = torch.from_numpy(arr.copy()) else: info_dict[k] = entry.list_data if info_dict: @@ -200,12 +209,11 @@ class OmniNPUModelRunner(NPUModelRunner): logger.error(f"Error decoding additional information: {e}") pass # ------------------------------------------------------------------------------------------------ - # Only relevant for models using M-RoPE (e.g, Qwen2-VL) if self.uses_mrope: - self._init_mrope_positions(self.requests[req_id]) + self._init_mrope_positions(req_state) - req_ids_to_add.append(req_id) + reqs_to_add.append(self.requests[req_id]) # Update the states of the running/resumed requests. is_last_rank = get_pp_group().is_last_rank @@ -250,7 +258,7 @@ class OmniNPUModelRunner(NPUModelRunner): # The request is not in the persistent batch. # The request was either preempted and resumed later, or was not # scheduled in the previous step and needs to be added again. - req_ids_to_add.append(req_id) + reqs_to_add.append(req_state) continue # Update the persistent batch. @@ -280,9 +288,8 @@ class OmniNPUModelRunner(NPUModelRunner): # Add the new or resumed requests to the persistent batch. # The smaller empty indices are filled first. - for req_id in req_ids_to_add: - req_state = self.requests[req_id] - self.input_batch.add_request(req_state) + for request in reqs_to_add: + self.input_batch.add_request(request) # Condense the batched states if there are gaps left by removed requests self.input_batch.condense() @@ -310,69 +317,11 @@ class OmniNPUModelRunner(NPUModelRunner): raise ValueError(f"Invalid hidden states type: {type(hidden_states)}") return text_hidden_states, multimodal_outputs - def _init_model_kwargs(self, num_tokens: int): - model_kwargs = dict[str, Any]() - - if not self.is_pooling_model: - return model_kwargs - - num_reqs = self.input_batch.num_reqs - pooling_params = self.input_batch.get_pooling_params() - - token_type_id_requests = dict[int, Any]() - for i, param in enumerate(pooling_params): - if ( - param.extra_kwargs is not None - and (token_types := param.extra_kwargs.get("compressed_token_type_ids")) is not None - ): - token_type_id_requests[i] = token_types - - if len(token_type_id_requests) == 0: - return model_kwargs - - seq_lens = self.seq_lens_cpu[:num_reqs] - token_type_ids = [] - - for i in range(num_reqs): - pos = token_type_id_requests.get(i, seq_lens[i]) - ids = (torch.arange(seq_lens[i]) >= pos).int() - token_type_ids.append(ids) - - model_kwargs["token_type_ids"] = torch.concat(token_type_ids).to(device=self.device) - return model_kwargs - - def _extract_mm_kwargs( - self, - scheduler_output: "SchedulerOutput", - ) -> BatchedTensorInputs: - if not scheduler_output or not self.is_multimodal_raw_input_only_model: - return {} - - mm_kwargs = list[MultiModalKwargsItem]() - for req in scheduler_output.scheduled_new_reqs: - for feature in req.mm_features: - if feature.data is not None: - mm_kwargs.append(feature.data) - - # Input all modalities at once - model = cast(SupportsMultiModal, self.model) - mm_kwargs_combined: BatchedTensorInputs = {} - for _, _, mm_kwargs_group in group_mm_kwargs_by_modality( - mm_kwargs, - device=self.device, - pin_memory=self.pin_memory, - merge_by_field_config=model.merge_by_field_config, - ): - mm_kwargs_combined.update(mm_kwargs_group) - - return mm_kwargs_combined - @torch.inference_mode() def _dummy_run( self, num_tokens: int, with_prefill: bool = False, - is_torchair_compile: bool = False, aclgraph_runtime_mode: CUDAGraphMode | None = None, force_attention: bool = False, uniform_decode: bool = False, @@ -383,7 +332,6 @@ class OmniNPUModelRunner(NPUModelRunner): CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL, } - # In multi-DP scenarios, there may be situations where all DP groups are executing dummy runs. # If sequence parallelism is enabled, it is essential to ensure that num_tokens is divisible by tp_size. if self.use_aclgraph and enable_sp(self.vllm_config): @@ -395,11 +343,7 @@ class OmniNPUModelRunner(NPUModelRunner): with_prefill = True # Padding for DP - (num_tokens, num_tokens_across_dp, with_prefill, _) = self._sync_metadata_across_dp( - num_tokens, with_prefill, False - ) - - moe_comm_type = self._select_moe_comm_method(num_tokens, with_prefill) + (num_tokens, num_tokens_across_dp, with_prefill) = self._sync_metadata_across_dp(num_tokens, with_prefill) # If cudagraph_mode.decode_mode() == FULL and # cudagraph_mode.separate_routine(). This means that we are using @@ -438,76 +382,109 @@ class OmniNPUModelRunner(NPUModelRunner): assert sum(num_scheduled_tokens_list) == num_tokens assert len(num_scheduled_tokens_list) == num_reqs num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32) + num_sampled_tokens = np.ones(num_reqs, dtype=np.int32) if not self.in_profile_run and self.dynamic_eplb: self.eplb_updator.forward_before() - with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens): + has_lora = True if self.lora_config and self.compilation_config.cudagraph_specialize_lora else False + _ag_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch( + num_tokens=num_tokens, uniform_decode=uniform_decode, has_lora=has_lora + ) + + num_tokens_padded = batch_descriptor.num_tokens + num_reqs_padded = batch_descriptor.num_reqs if batch_descriptor.num_reqs is not None else num_reqs + if num_tokens_across_dp is not None and num_tokens_padded != num_tokens: + # pad is needed if the pad of `num_tokens` is triggered inside CudagraphDispatcher + num_tokens_across_dp[:] = num_tokens_padded + num_scheduled_tokens = num_scheduled_tokens.repeat(num_reqs_padded) + + moe_comm_type = self._select_moe_comm_method(num_tokens_padded) + + # filter out the valid batch descriptor + if aclgraph_runtime_mode is not None: + # we allow forcing NONE when the dispatcher disagrees to support + # warm ups for aclgraph capture + if aclgraph_runtime_mode != CUDAGraphMode.NONE and aclgraph_runtime_mode != _ag_mode: + raise ValueError( + f"Aclgraph runtime mode mismatch at dummy_run. " + f"Expected {_ag_mode}, but got {aclgraph_runtime_mode}." + ) + else: + aclgraph_runtime_mode = _ag_mode + + # TODO(Mengqing): Set create_mixed_batch to False since it's only used in FI warmup + # and not supported in ASCEND now. We could remove it in the future. + attn_metadata = self._build_dummy_attn_metadata( + False, + num_reqs=num_reqs_padded, + num_tokens=num_tokens_padded, + max_query_len=max_query_len, + aclgraph_runtime_mode=aclgraph_runtime_mode, + force_attention=force_attention, + num_scheduled_tokens=num_scheduled_tokens, + ) + + with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, num_sampled_tokens): + # Make sure padding doesn't exceed max_num_tokens + assert num_tokens_padded <= self.max_num_tokens if self.is_multimodal_model: input_ids = None - inputs_embeds = self.inputs_embeds[:num_tokens] + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] + elif self.enable_prompt_embeds: + input_ids = None + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] else: - input_ids = self.input_ids[:num_tokens] + input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None if self.uses_mrope: - positions = self.mrope_positions[:, :num_tokens] + positions = self.mrope_positions.gpu[:, :num_tokens_padded] else: - positions = self.positions[:num_tokens] + positions = self.positions.gpu[:num_tokens_padded] if get_pp_group().is_first_rank: intermediate_tensors = None else: + # When PP and flashcomm1 are enabled, + # during dummy_run the estimated space should divide num_tokens by tp_size; + # otherwise, on non-first PP ranks it would effectively perform an extra all-gather, + # leading to incorrect memory estimation and potentially causing OOM. + actual_tokens = num_tokens + if enable_sp(): + tp_size = get_tensor_model_parallel_world_size() + actual_tokens = num_tokens // tp_size if self.intermediate_tensors is None: self.intermediate_tensors = self.model.make_empty_intermediate_tensors( - batch_size=num_tokens, dtype=self.dtype, device=self.device + batch_size=actual_tokens, dtype=self.dtype, device=self.device ) intermediate_tensors = IntermediateTensors( - {k: v[:num_tokens] for k, v in self.intermediate_tensors.items()} + {k: v[:num_tokens_padded] for k, v in self.intermediate_tensors.items()} ) - # filter out the valid batch descriptor - _ag_mode, batch_descriptor = self.aclgraph_dispatcher.dispatch( - BatchDescriptor(num_tokens=num_tokens, uniform_decode=uniform_decode) - ) - if aclgraph_runtime_mode is not None: - # we allow forcing NONE when the dispatcher disagrees to support - # warm ups for aclgraph capture - assert aclgraph_runtime_mode == CUDAGraphMode.NONE or aclgraph_runtime_mode == _ag_mode, ( - f"Aclgraph runtime mode mismatch at dummy_run. " - f"Expected {_ag_mode}, but got {aclgraph_runtime_mode}." - ) - else: - aclgraph_runtime_mode = _ag_mode - - # TODO(Mengqing): Set create_mixed_batch to False since it's only used in FI warmup - # and not supported in ASCEND now. We could remove it in the future. - attn_metadata = self._build_dummy_attn_metadata( - False, - num_reqs=num_reqs, - num_tokens=num_tokens, - max_query_len=max_query_len, - aclgraph_runtime_mode=aclgraph_runtime_mode, - force_attention=force_attention, - ) - need_dummy_logits = not self.in_profile_run and lmhead_tp_enable() + max_num_reqs_across_dp = max_num_reqs * self.uniform_decode_query_len + dummy_indices = torch.zeros(max_num_reqs_across_dp, dtype=torch.int32) - if need_dummy_logits: - max_num_reqs_across_dp = num_tokens if not with_prefill else max_num_reqs - dummy_indices = torch.zeros(max_num_reqs_across_dp, dtype=torch.int32) + def dummy_compute_logits(hidden_states): + if not need_dummy_logits: + return None + return self.model.compute_logits(hidden_states[dummy_indices]) - def dummy_compute_logits(hidden_states): - return self.model.compute_logits(hidden_states[dummy_indices]) + def dummy_drafter_compute_logits(hidden_states): + if not need_dummy_logits or self.drafter is None: + return + if hasattr(self.drafter, "model") and hasattr(self.drafter.model, "compute_logits"): + return self.drafter.model.compute_logits(hidden_states[dummy_indices]) with set_ascend_forward_context( attn_metadata, self.vllm_config, - num_tokens=num_tokens, + num_tokens=num_tokens_padded, num_tokens_across_dp=num_tokens_across_dp, with_prefill=with_prefill, in_profile_run=self.in_profile_run, - reserved_mc2_mask=self.reserved_mc2_mask, + # reserved_mc2_mask=self.reserved_mc2_mask, moe_comm_type=moe_comm_type, num_actual_tokens=0, aclgraph_runtime_mode=aclgraph_runtime_mode, @@ -517,30 +494,21 @@ class OmniNPUModelRunner(NPUModelRunner): weight_prefetch_method=self.weight_prefetch_method, ): hidden_states = self._generate_dummy_run_hidden_states( - with_prefill, - is_torchair_compile, - input_ids, - positions, - attn_metadata, - num_tokens, - intermediate_tensors, - inputs_embeds, + input_ids, positions, num_tokens_padded, intermediate_tensors, inputs_embeds ) - if need_dummy_logits: - dummy_compute_logits(hidden_states) + dummy_compute_logits(hidden_states) if self.drafter: self.drafter.dummy_run( - num_tokens=num_tokens, + num_tokens=num_tokens_padded, with_prefill=with_prefill, - skip_attn=True, - num_reqs=num_reqs, + num_reqs=num_reqs_padded, num_tokens_across_dp=num_tokens_across_dp, aclgraph_runtime_mode=aclgraph_runtime_mode, batch_descriptor=batch_descriptor, + dummy_compute_logits=dummy_drafter_compute_logits, + skip_attn=not force_attention, ) - if need_dummy_logits: - self.drafter.model.compute_logits(hidden_states[dummy_indices]) if self.in_profile_run and self.dynamic_eplb: self.model.clear_all_moe_loads() if not self.in_profile_run and self.dynamic_eplb: @@ -550,3 +518,119 @@ class OmniNPUModelRunner(NPUModelRunner): hidden_states, _ = self.extract_multimodal_outputs(hidden_states) return hidden_states + + def _decode_and_store_request_payloads(self, scheduler_output: "SchedulerOutput") -> None: + """Decode per-request prompt_embeds and additional_information for newly + scheduled requests and store them to CPU in the request state. + This version avoids hard dependency on payload classes by duck-typing.""" + try: + new_reqs = getattr(scheduler_output, "scheduled_new_reqs", []) + if not new_reqs: + return + for nr in new_reqs: + req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None) + if req_id is None: + continue + # prompt_embeds + payload_pe = getattr(nr, "prompt_embeds", None) + pe_cpu = None + if payload_pe is not None: + if isinstance(payload_pe, torch.Tensor): + pe_cpu = payload_pe.detach().to("cpu").contiguous() + else: + # Try duck-typing a payload with data/shape/dtype + data = getattr(payload_pe, "data", None) + shape = getattr(payload_pe, "shape", None) + if data is not None and shape is not None: + dt = np.dtype(getattr(payload_pe, "dtype", "float32")) + arr = np.frombuffer(data, dtype=dt) + arr = arr.reshape(shape) + pe_cpu = torch.from_numpy(arr.copy()) + if pe_cpu is not None and req_id in self.requests: + setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) + # additional_information + payload_info = getattr(nr, "additional_information", None) + if payload_info is not None: + info_dict = {} + if isinstance(payload_info, dict): + info_dict = payload_info + else: + # Try duck-typing a payload with entries, each entry may have + # tensor_data/tensor_dtype/tensor_shape or list_data + entries = getattr(payload_info, "entries", None) + if isinstance(entries, dict): + for k, entry in entries.items(): + tensor_data = getattr(entry, "tensor_data", None) + if tensor_data is not None: + dt = np.dtype(getattr(entry, "tensor_dtype", "float32")) + arr = np.frombuffer(tensor_data, dtype=dt) + arr = arr.reshape(getattr(entry, "tensor_shape", ())) + info_dict[k] = torch.from_numpy(arr.copy()) + else: + info_dict[k] = getattr(entry, "list_data", None) + if info_dict and req_id in self.requests: + setattr(self.requests[req_id], "additional_information_cpu", info_dict) + except Exception as e: + logger.error(f"Error decoding prompt_embeds / additional_information: {e}") + + def _gather_runtime_additional_information(self) -> list[dict]: + """Gather per-request additional_information stored in request state in batch order.""" + per_req_runtime_info = [] + for req_id in self.input_batch.req_ids: + req_state = self.requests.get(req_id) + info = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None + if info and isinstance(info, dict): + per_req_runtime_info.append(info) + if "thinker_reply_part_per_request" in info: + q = info["thinker_reply_part_per_request"] + if hasattr(q, "shape"): + logger.debug(f"[OMNI] req={req_id} has thinker_reply_part_per_request queue shape: {q.shape}") + else: + per_req_runtime_info.append({}) + return per_req_runtime_info + + def _compute_request_token_spans(self, num_scheduled_tokens_np) -> list[tuple[int, int]]: + """Compute (start, end) token spans for each request within the flattened step sequence.""" + req_token_spans: list[tuple[int, int]] = [] + for req_index in range(len(self.input_batch.req_ids)): + start_offset = int(self.query_start_loc.cpu[req_index]) + sched_tokens = int(num_scheduled_tokens_np[req_index]) + req_token_spans.append((start_offset, start_offset + sched_tokens)) + return req_token_spans + + def _build_model_kwargs_extra(self) -> dict: + """Build extra keyword arguments passed to the model for this step, including: + - runtime_additional_information: per-request additional information stored in request state + """ + model_kwargs_extra: dict[str, object] = {} + try: + model_kwargs_extra["runtime_additional_information"] = self._gather_runtime_additional_information() + except Exception as e: + logger.error(f"[OMNI DEBUG] Error building model_kwargs_extra: {e}") + import traceback + + traceback.print_exc() + return model_kwargs_extra + + def _collect_additional_information_for_prefill( + self, + num_scheduled_tokens_np: np.ndarray, + ) -> dict[str, dict]: + """Overlay per-request prompt_embeds for the prefill portion and collect + additional_information slices for this step. Returns a map req_id -> dict.""" + for req_index, req_id in enumerate(self.input_batch.req_ids): + req_state = self.requests[req_id] + pe_cpu = getattr(req_state, "prompt_embeds_cpu", None) + num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index]) + prompt_len = len(req_state.prompt_token_ids) + prompt_remaining = max(0, prompt_len - num_computed_tokens) + sched_tokens = int(num_scheduled_tokens_np[req_index]) + overlay_len = min(sched_tokens, prompt_remaining) + if overlay_len <= 0: + continue + if overlay_len > 0 and pe_cpu is not None: + src = pe_cpu[num_computed_tokens : num_computed_tokens + overlay_len].to( + dtype=self.dtype, device=self.device, non_blocking=True + ) + start_offset = int(self.query_start_loc.cpu[req_index]) + self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src)