diff --git a/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml b/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml index a32ca7d2d..7ce88f80a 100644 --- a/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml +++ b/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml @@ -20,7 +20,7 @@ stage_args: trust_remote_code: true engine_output_type: text distributed_executor_backend: "mp" - enable_prefix_caching: false + enable_prefix_caching: true hf_config_name: thinker_config tensor_parallel_size: 2 final_output: true diff --git a/vllm_omni/core/prefix_cache.py b/vllm_omni/core/prefix_cache.py index 25acfe8ed..17a981140 100644 --- a/vllm_omni/core/prefix_cache.py +++ b/vllm_omni/core/prefix_cache.py @@ -2,6 +2,8 @@ Utilities for Prefix Caching in Omni models. """ +from dataclasses import dataclass, field + import torch from vllm.logger import init_logger from vllm.v1.worker.gpu_input_batch import InputBatch @@ -11,6 +13,23 @@ from vllm_omni.utils.mm_outputs import build_mm_cpu, to_payload_element logger = init_logger(__name__) +@dataclass +class _PendingAsyncWrite: + """One in-flight async write into the prefix cache. + + Holds the CPU tensors that PyTorch implicitly allocated via + ``.to("cpu", non_blocking=True)`` together with a CUDA event recorded + on the copy stream. The consumer calls ``event.synchronize()`` once + (blocking) to guarantee the D2H has completed before reading. + """ + + event: "torch.cuda.Event" + num_tokens: int + slots_cpu: torch.Tensor + hidden_cpu: torch.Tensor | None + mm_cpu: dict[str, torch.Tensor] = field(default_factory=dict) + + class OmniTensorPrefixCache: """Prefix cache for hidden states (model outputs) and model specific multimodal outputs. @@ -49,6 +68,19 @@ class OmniTensorPrefixCache: self._new_req_cache_hit_ids: set[str] = set() self._deferred_mm_outputs: dict[str, dict[str, list[tuple[int, torch.Tensor]]]] = {} + # Async write pipeline state: + # - one dedicated copy stream + # - one in-flight pending write at a time (the producer is exactly + # one forward step ahead of the consumer) + # - schedule_async_write() consumes the *previous* pending write + # (via event.synchronize() + scatter) before kicking off the new + # D2H. By the time we hit synchronize, the previous step's GPU + # work has been done for ~one forward step, so the sync is + # near-zero wall-clock. + self._async_initialized: bool = False + self._async_copy_stream: torch.cuda.Stream | None = None + self._pending_write: _PendingAsyncWrite | None = None + def maybe_init_missing_mm_cache_keys(self, multimodal_outputs: dict, seq_len: int): """Given multimodal outputs from executing the model, dynamically determine which multimodal outputs are tensors depending on sequence @@ -84,14 +116,186 @@ class OmniTensorPrefixCache: logger.info("Initializing multimodal output cache of size %s for key: %s", list(new_tensor_shape), key) def _get_cache_tensor(self, dtype: torch.dtype, hidden_size: int | None = None) -> torch.Tensor: - """Allocate a CPU cache tensor for a specific key.""" + """Allocate a CPU cache tensor for a specific key. + + When CUDA is available the tensor is pinned: this is what lets the + async-write pipeline ride the GPU->CPU copy on the dedicated copy + stream as a true async ``cudaMemcpyAsync`` and lets the post-copy CPU + scatter avoid page-resolution stalls. On CPU-only builds pinning is + unsupported (``torch.zeros(pin_memory=True)`` raises), and the async + pipeline is disabled anyway (``schedule_async_write`` early-returns + when ``torch.cuda.is_available()`` is False), so we fall back to a + plain pageable allocation with no behavioral change. + """ actual_hidden_size = hidden_size if hidden_size is not None else self.default_hidden_size return torch.zeros( (self.num_blocks, self.block_size, actual_hidden_size), dtype=dtype, device="cpu", + pin_memory=torch.cuda.is_available(), ) + # ===================================================================== + # Async write pipeline: + # schedule_async_write : called per forward step, enqueues a + # non-blocking D2H + records a CUDA event + # drain_ready_async_writes : non-blocking poll, applies any writes + # whose D2H has completed + # ===================================================================== + + def _ensure_async_resources(self) -> None: + """Lazily allocate the dedicated copy stream. + + The per-step CPU staging tensors are allocated by PyTorch on demand + inside ``.to("cpu", non_blocking=True)``, which gives them + implicitly-pinned memory and lets PyTorch's allocator amortize / + reuse them. + """ + if self._async_initialized: + return + if not torch.cuda.is_available(): + return + self._async_copy_stream = torch.cuda.Stream() + self._async_initialized = True + logger.info("Omni prefix cache async-write copy stream initialized.") + + def schedule_async_write( + self, + hidden_states_gpu: torch.Tensor | None, + multimodal_outputs_gpu: dict[str, torch.Tensor] | None, + slot_mapping_gpu: torch.Tensor, + num_tokens_unpadded: int, + num_tokens_padded: int, + skip_mm_cache_keys: set[str] | None = None, + ) -> None: + """Kick off a non-blocking GPU->CPU copy for this step's writes. + + 1. **Consume the previous step's write** (if any): block on its + event and scatter the CPU tensors into the cache. By the time we + reach this point, the previous step's GPU work has been done for + ~one forward step, so the synchronize is near-zero wall-clock. + 2. **Schedule THIS step's D2H** on a dedicated copy stream using + ``.to("cpu", non_blocking=True)``. PyTorch allocates the pinned + destination implicitly. + 3. **Record an event** and stash everything for the next step's + consume. + + Holding exactly one in-flight write at a time keeps the design + simple: no ring, no event-polling, no driver-mutex pressure from + repeated event.query() calls. + """ + if not torch.cuda.is_available(): + return + + self._ensure_async_resources() + if not self._async_initialized: + return + + # Decide which mm keys we'll cache. Init the dest cache tensors + # before scheduling the copies so the apply step has somewhere + # to scatter into. + skip = skip_mm_cache_keys or set() + cacheable_mm_keys: list[str] = [] + if multimodal_outputs_gpu: + self.maybe_init_missing_mm_cache_keys(multimodal_outputs_gpu, seq_len=num_tokens_padded) + for k in self.mm_cache_keys: + if k in skip: + continue + v = multimodal_outputs_gpu.get(k) + if not isinstance(v, torch.Tensor) or v.ndim < 2: + continue + if v.shape[0] < num_tokens_unpadded: + continue + cacheable_mm_keys.append(k) + + # Step 1: consume the previous step's write FIRST. The previous + # step's event was recorded ~one forward step ago, so + # synchronize() is effectively a no-op cost-wise; the scatter + # into pinned CPU caches is also fast. + self._consume_pending_write() + + if hidden_states_gpu is None and not cacheable_mm_keys: + return # nothing to schedule this step + + # Step 2: schedule the D2H on the copy stream. PyTorch implicitly + # allocates pinned host destinations for the ``.to("cpu", + # non_blocking=True)`` results. + copy_stream = self._async_copy_stream + default_stream = torch.cuda.current_stream() + with torch.cuda.stream(copy_stream): + copy_stream.wait_stream(default_stream) + + # slot_mapping: copy as-is; the int64 cast required by + # ``index_copy_`` is deferred to ``_consume_pending_write`` + # so we avoid launching an extra GPU kernel here that would + # contend the driver mutex with the compute stream. + slots_cpu = slot_mapping_gpu[:num_tokens_unpadded].to("cpu", non_blocking=True) + hidden_cpu: torch.Tensor | None = None + if hidden_states_gpu is not None: + hidden_cpu = hidden_states_gpu[:num_tokens_unpadded].to("cpu", non_blocking=True) + mm_cpu: dict[str, torch.Tensor] = {} + for k in cacheable_mm_keys: + mm_cpu[k] = multimodal_outputs_gpu[k][:num_tokens_unpadded].to("cpu", non_blocking=True) + + event = torch.cuda.Event() + event.record() + + # Step 3: stash for next step's consume. + self._pending_write = _PendingAsyncWrite( + event=event, + num_tokens=num_tokens_unpadded, + slots_cpu=slots_cpu, + hidden_cpu=hidden_cpu, + mm_cpu=mm_cpu, + ) + + def drain_ready_async_writes(self) -> int: + """Consume one in-flight write if it has finished (non-blocking). + + Exposed so the model runner can drain at the top of + ``execute_model`` regardless of whether a new write will be + scheduled this step. + """ + if self._pending_write is None: + return 0 + if not self._pending_write.event.query(): + return 0 + self._consume_pending_write() + return 1 + + def _consume_pending_write(self) -> None: + """Synchronize on the pending event and scatter the CPU tensors + into the destination caches. Drops state on completion. + """ + entry = self._pending_write + if entry is None: + return + self._pending_write = None + entry.event.synchronize() + n = entry.num_tokens + if n <= 0: + return + + slots = entry.slots_cpu + if slots.dtype != torch.int64: + slots = slots.to(torch.int64) + + # NOTE: ``index_copy_(0, idx, src)`` is the row-scatter primitive in + # PyTorch. It's semantically identical to ``flat[idx] = src`` but + # dispatches to a much simpler CPU implementation (single dim, + # no broadcasting, no accumulate), which on pinned memory at large + # row sizes is observed to be much faster than the generic + # ``aten::index_put_`` path used by advanced-indexing assignment. + if entry.hidden_cpu is not None: + flat = self.hidden_states_cache.view(-1, self.hidden_states_cache.shape[-1]) + flat.index_copy_(0, slots, entry.hidden_cpu) + for k, src_cpu in entry.mm_cpu.items(): + mm_cache = self.mm_outputs_cache.get(k) + if mm_cache is None: + continue + flat = mm_cache.view(-1, mm_cache.shape[-1]) + flat.index_copy_(0, slots, src_cpu) + def add_prefix_cached_new_req_id(self, req_id: str): """Adds a new request ID to the set of prefix cache hits on the batch.""" self._new_req_cache_hit_ids.add(req_id) @@ -173,9 +377,17 @@ class OmniTensorPrefixCache: num_tokens_unpadded, hidden_states_cpu, ) - # View the cache as 2D so that we can treat our slots as row indices + # View the cache as 2D so that we can treat our slots as row indices. + # ``index_copy_`` is the row-scatter primitive — it's semantically + # ``cache[idx] = src`` but dispatches to a faster CPU path than the + # generic advanced-indexing assignment (which goes through + # ``aten::index_put_`` and is dominated by per-row dispatch + # overhead on large feature widths). flat_cache = self.hidden_states_cache.view(-1, self.hidden_states_cache.shape[-1]) - flat_cache[unpadded_slot_mapping] = hidden_states + slot_idx = unpadded_slot_mapping + if slot_idx.dtype != torch.int64: + slot_idx = slot_idx.to(torch.int64) + flat_cache.index_copy_(0, slot_idx, hidden_states) logger.debug("Writing to hidden states for %s tokens", num_tokens_unpadded) # Do the same for the stage's cached multimodal outputs @@ -195,7 +407,10 @@ class OmniTensorPrefixCache: mm_state = multimodal_outputs[mm_out_key][:num_tokens_unpadded] mm_state = OmniTensorPrefixCache._coerce_to_cpu_tensor(mm_state) flat_cache = mm_cache.view(-1, mm_cache.shape[-1]) - flat_cache[unpadded_slot_mapping] = mm_state + slot_idx = unpadded_slot_mapping + if slot_idx.dtype != torch.int64: + slot_idx = slot_idx.to(torch.int64) + flat_cache.index_copy_(0, slot_idx, mm_state) logger.debug("Writing to mm output cache for %s tokens", num_tokens_unpadded) def stage_deferred_mm_outputs( @@ -331,7 +546,9 @@ class OmniTensorPrefixCache: mm_state_cpu = OmniTensorPrefixCache._coerce_to_cpu_tensor(mm_state[: slots.numel()]) mm_cache = self.mm_outputs_cache[mm_key] flat_cache = mm_cache.view(-1, mm_cache.shape[-1]) - flat_cache[slots] = mm_state_cpu + if slots.dtype != torch.int64: + slots = slots.to(torch.int64) + flat_cache.index_copy_(0, slots, mm_state_cpu) def _get_slot_ids_for_token_range( self, diff --git a/vllm_omni/deploy/qwen3_omni_moe.yaml b/vllm_omni/deploy/qwen3_omni_moe.yaml index 06869a657..e807e13c7 100644 --- a/vllm_omni/deploy/qwen3_omni_moe.yaml +++ b/vllm_omni/deploy/qwen3_omni_moe.yaml @@ -28,7 +28,7 @@ stages: max_num_seqs: 64 gpu_memory_utilization: 0.9 trust_remote_code: true - enable_prefix_caching: false + enable_prefix_caching: true devices: "0" default_sampling_params: temperature: 0.0 diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index e55efd3a1..b1f1805fb 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -500,6 +500,16 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): if hasattr(self.model, "_clear_warmup_state"): self.model._clear_warmup_state() + # Async-write pipeline: apply any pending GPU->CPU prefix-cache writes + # whose copy event has already fired. Non-blocking — entries whose D2H + # is still in flight stay queued and will be picked up on the next + # step's drain. This guarantees any downstream read of + # ``omni_prefix_cache.hidden_states_cache`` / + # ``omni_prefix_cache.mm_outputs_cache`` in this step sees the + # state produced no later than the previous forward step. + if self.omni_prefix_cache is not None: + self.omni_prefix_cache.drain_ready_async_writes() + # [Omni] Handle KV transfer BEFORE updating states (which removes finished requests) finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}) if finished_reqs and hasattr(self.model, "get_kv_transfer_metadata"): @@ -796,18 +806,30 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output) hidden_states_cpu = None - if self.omni_prefix_cache is not None and get_pp_group().is_last_rank: - hidden_states_cpu = hidden_states[:num_tokens_unpadded].detach().to("cpu").contiguous() - # Cache hidden states & multimodal outputs if we've enabled hidden state - # prefix caching unless this isn't the last pipeline parallelism rank. - self._maybe_update_prefix_cache( - hidden_states=hidden_states, - hidden_states_cpu=hidden_states_cpu, - multimodal_outputs=multimodal_outputs, - num_tokens_unpadded=num_tokens_unpadded, - num_tokens_padded=num_tokens_padded, - ) + # Async-write pipeline (replaces the per-step blocking + # ``.to("cpu")`` + ``aten::index_put_`` on pageable host memory). + # Schedules non-blocking GPU->CPU copies on a dedicated stream; + # the actual CPU scatter into ``hidden_states_cache`` / + # ``mm_outputs_cache`` happens in ``drain_ready_async_writes`` + # at the top of subsequent execute_model() calls. + if self.omni_prefix_cache is not None and get_pp_group().is_last_rank: + hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None + # Some models (e.g. qwen3-tts-talker) opt out of full-hidden-state + # prefix caching but the downstream pooler payload path still + # needs a CPU hidden-states view. Materialize it synchronously + # in that case; the legacy behavior is preserved. + if hs_for_cache is None: + hidden_states_cpu = hidden_states[:num_tokens_unpadded].detach().to("cpu").contiguous() + slot_mapping_gpu = self.input_batch.block_table[0].slot_mapping.gpu + self.omni_prefix_cache.schedule_async_write( + hidden_states_gpu=hs_for_cache, + multimodal_outputs_gpu=(flatten_payload(multimodal_outputs) if multimodal_outputs else None), + slot_mapping_gpu=slot_mapping_gpu, + num_tokens_unpadded=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded, + skip_mm_cache_keys=self._deferred_prefix_cache_mm_keys(), + ) if not self.broadcast_pp_output: # Common case.