diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index 0b39a4000..88871d55f 100644 --- a/vllm/_xpu_ops.py +++ b/vllm/_xpu_ops.py @@ -1,16 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING +import json +import os +import re +import time +from typing import TYPE_CHECKING, Any import torch from vllm_xpu_kernels.flash_attn_interface import flash_attn_varlen_func from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.utils.xpu_decode_timing import timed_region from vllm.utils.torch_utils import direct_register_custom_op logger = init_logger(__name__) +_XPU_GDN_TRACE_LINES = 0 if TYPE_CHECKING: @@ -92,6 +98,272 @@ if hasattr(torch.ops._xpu_C, "int4_gemm_w4a16"): return torch.empty((M, N), dtype=input.dtype, device=input.device) +def _xpu_gdn_trace_int_env(name: str, default: int) -> int: + try: + return int(os.environ.get(name, str(default)) or default) + except ValueError: + return default + + +def _xpu_gdn_trace_rank() -> int: + try: + from vllm.distributed import get_tensor_model_parallel_rank + + return int(get_tensor_model_parallel_rank()) + except Exception: + try: + return int(os.environ.get("LOCAL_RANK", "0") or "0") + except ValueError: + return 0 + + +def _xpu_gdn_json_number(value: Any) -> int | float | str | None: + try: + if value is None: + return None + if isinstance(value, bool): + return bool(value) + if isinstance(value, int): + return int(value) + if isinstance(value, float): + if value == float("inf"): + return "inf" + if value == float("-inf"): + return "-inf" + if value != value: + return "nan" + return value + except Exception: + return repr(value) + return value + + +def _xpu_gdn_tensor_digest( + tensor: torch.Tensor | None, + *, + head_limit: int, + include_head: bool = False, +) -> dict[str, Any] | None: + if tensor is None: + return None + try: + with torch.no_grad(): + t = tensor.detach() + record: dict[str, Any] = { + "shape": list(t.shape), + "dtype": str(t.dtype), + "device": str(t.device), + "numel": int(t.numel()), + } + if t.numel() == 0: + record.update({"finite": 0, "nan": 0, "posinf": 0, "neginf": 0}) + return record + stats_t = t.float() if not t.is_floating_point() else t + finite_mask = torch.isfinite(stats_t) + nan_mask = torch.isnan(stats_t) + posinf_mask = stats_t == float("inf") + neginf_mask = stats_t == float("-inf") + finite_count = int(finite_mask.sum().item()) + record.update( + { + "finite": finite_count, + "nan": int(nan_mask.sum().item()), + "posinf": int(posinf_mask.sum().item()), + "neginf": int(neginf_mask.sum().item()), + } + ) + if finite_count: + finite_vals = stats_t[finite_mask].float() + record.update( + { + "min": _xpu_gdn_json_number(finite_vals.min().item()), + "max": _xpu_gdn_json_number(finite_vals.max().item()), + "mean": _xpu_gdn_json_number(finite_vals.mean().item()), + "sum": _xpu_gdn_json_number(finite_vals.sum().item()), + "l2": _xpu_gdn_json_number( + torch.linalg.vector_norm(finite_vals).item() + ), + } + ) + if include_head and head_limit > 0: + head = t.reshape(-1)[:head_limit].to("cpu").tolist() + record["head"] = [_xpu_gdn_json_number(v) for v in head] + return record + except Exception as exc: + return {"error": repr(exc)} + + +def _xpu_gdn_selected_state_digest( + state: torch.Tensor | None, + indices: torch.Tensor | None, + *, + head_limit: int, + state_limit: int, +) -> dict[str, Any] | None: + if state is None or indices is None: + return None + try: + flat_indices = indices.detach().reshape(-1) + valid = flat_indices[(flat_indices >= 0) & (flat_indices < state.shape[0])] + if state_limit > 0: + valid = valid[:state_limit] + record: dict[str, Any] = { + "indices_shape": list(indices.shape), + "indices_head": [ + int(v) for v in flat_indices[: max(0, head_limit)].to("cpu").tolist() + ], + "selected_count": int(valid.numel()), + } + if valid.numel() > 0: + selected = state.index_select(0, valid.to(device=state.device).long()) + record["selected"] = _xpu_gdn_tensor_digest( + selected, head_limit=head_limit, include_head=True + ) + return record + except Exception as exc: + return {"error": repr(exc)} + + +def _xpu_gdn_trace_enabled( + forward_context: Any, + layer_name: str, + attn_metadata: Any, +) -> tuple[bool, dict[str, Any]]: + trace_file = os.environ.get("VLLM_XPU_GDN_TRACE_FILE", "") + if not trace_file: + return False, {} + + global _XPU_GDN_TRACE_LINES + max_lines = _xpu_gdn_trace_int_env("VLLM_XPU_GDN_TRACE_MAX_LINES", 0) + if max_lines > 0 and _XPU_GDN_TRACE_LINES >= max_lines: + return False, {} + + tp_rank = _xpu_gdn_trace_rank() + rank_filter = os.environ.get("VLLM_XPU_GDN_TRACE_RANK", "") + if rank_filter and rank_filter not in ("*", "all") and str(tp_rank) != rank_filter: + return False, {} + + layer_regex = os.environ.get("VLLM_XPU_GDN_TRACE_LAYER_REGEX", "") + if layer_regex: + try: + if not re.search(layer_regex, layer_name): + return False, {} + except re.error: + return False, {} + + if ( + os.environ.get("VLLM_XPU_GDN_TRACE_PREFILL_ONLY", "0") == "1" + and getattr(attn_metadata, "num_prefills", 0) <= 0 + ): + return False, {} + if ( + os.environ.get("VLLM_XPU_GDN_TRACE_DECODE_ONLY", "0") == "1" + and getattr(attn_metadata, "num_decodes", 0) <= 0 + ): + return False, {} + + extra = getattr(forward_context, "additional_kwargs", {}) or {} + req_ids = [str(req_id) for req_id in extra.get("xpu_req_ids", [])] + req_regex = os.environ.get("VLLM_XPU_GDN_TRACE_REQ_REGEX", "") + matched_req_ids = req_ids + if req_regex: + try: + matched_req_ids = [req_id for req_id in req_ids if re.search(req_regex, req_id)] + except re.error: + matched_req_ids = [] + if not matched_req_ids: + return False, {} + + return True, { + "trace_file": trace_file, + "tp_rank": tp_rank, + "req_ids": req_ids, + "matched_req_ids": matched_req_ids, + "forward_context_extra": extra, + } + + +def _xpu_gdn_trace_write( + *, + stage: str, + trace_context: dict[str, Any], + layer_name: str, + attn_metadata: Any, + core_attn_out: torch.Tensor | None, + z: torch.Tensor | None, + projected_states_qkvz: torch.Tensor | None, + projected_states_ba: torch.Tensor | None, + conv_state: torch.Tensor | None, + ssm_state: torch.Tensor | None, + non_spec_state_indices_tensor: torch.Tensor | None, +) -> None: + global _XPU_GDN_TRACE_LINES + max_lines = _xpu_gdn_trace_int_env("VLLM_XPU_GDN_TRACE_MAX_LINES", 0) + if max_lines > 0 and _XPU_GDN_TRACE_LINES >= max_lines: + return + head_limit = max(0, _xpu_gdn_trace_int_env("VLLM_XPU_GDN_TRACE_TENSOR_LIMIT", 8)) + state_limit = max(1, _xpu_gdn_trace_int_env("VLLM_XPU_GDN_TRACE_STATE_LIMIT", 2)) + try: + has_initial_state = getattr(attn_metadata, "has_initial_state", None) + record = { + "ts": time.time(), + "pid": os.getpid(), + "stage": stage, + "tp_rank": trace_context.get("tp_rank"), + "layer": layer_name, + "req_ids": trace_context.get("req_ids", []), + "matched_req_ids": trace_context.get("matched_req_ids", []), + "forward_context_extra": trace_context.get("forward_context_extra", {}), + "metadata": { + "num_prefills": int(getattr(attn_metadata, "num_prefills", 0)), + "num_decodes": int(getattr(attn_metadata, "num_decodes", 0)), + "num_actual_tokens": int( + getattr(attn_metadata, "num_actual_tokens", 0) + ), + "has_initial_state": None + if has_initial_state is None + else [ + bool(v) + for v in has_initial_state.detach() + .reshape(-1)[:head_limit] + .to("cpu") + .tolist() + ], + }, + "projected_states_qkvz": _xpu_gdn_tensor_digest( + projected_states_qkvz, head_limit=head_limit, include_head=True + ), + "projected_states_ba": _xpu_gdn_tensor_digest( + projected_states_ba, head_limit=head_limit, include_head=True + ), + "core_attn_out": _xpu_gdn_tensor_digest( + core_attn_out, head_limit=head_limit, include_head=True + ), + "z": _xpu_gdn_tensor_digest(z, head_limit=head_limit, include_head=True), + "conv_state": _xpu_gdn_selected_state_digest( + conv_state, + non_spec_state_indices_tensor, + head_limit=head_limit, + state_limit=state_limit, + ), + "ssm_state": _xpu_gdn_selected_state_digest( + ssm_state, + non_spec_state_indices_tensor, + head_limit=head_limit, + state_limit=state_limit, + ), + } + trace_file = str(trace_context["trace_file"]) + trace_dir = os.path.dirname(trace_file) + if trace_dir: + os.makedirs(trace_dir, exist_ok=True) + with open(trace_file, "a", encoding="utf-8") as f: + f.write(json.dumps(record, separators=(",", ":")) + "\n") + _XPU_GDN_TRACE_LINES += 1 + except Exception: + logger.exception("Failed to write XPU GDN trace") + + def _gdn_attention_core_xpu_impl( core_attn_out: torch.Tensor, z: torch.Tensor, @@ -114,38 +386,171 @@ def _gdn_attention_core_xpu_impl( attn_metadata = attn_metadata_raw[self.prefix] assert isinstance(attn_metadata, GDNAttentionMetadata) - # TODO: xpu does not support speculative decoding yet - assert attn_metadata.spec_sequence_masks is None # type: ignore[attr-defined] + def unpack_projected_states() -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor + ]: + if self.gqa_interleaved_layout: + from einops import rearrange + + query, key, value, z_generic, b, a = self.fix_query_key_value_ordering( + projected_states_qkvz, projected_states_ba + ) + query, key, value = map( + lambda x: rearrange(x, "l p d -> l (p d)"), (query, key, value) + ) + mixed_qkv = torch.cat((query, key, value), dim=-1) + else: + qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size + z_size = self.value_dim // self.tp_size + mixed_qkv, z_generic = projected_states_qkvz.split( + [qkv_size, z_size], dim=-1 + ) + z_generic = z_generic.reshape(z_generic.size(0), -1, self.head_v_dim) + b, a = projected_states_ba.chunk(2, dim=-1) + b = b.contiguous() + a = a.contiguous() + return mixed_qkv, z_generic, b, a + + gdn_native_fallback = os.environ.get( + "VLLM_XPU_GDN_NATIVE_FALLBACK", "" + ).strip().lower() + if gdn_native_fallback in ("decode", "decode_only") and ( + attn_metadata.spec_sequence_masks is None # type: ignore[attr-defined] + and attn_metadata.num_prefills == 0 # type: ignore[attr-defined] + and attn_metadata.num_decodes > 0 # type: ignore[attr-defined] + ): + mixed_qkv, z_generic, b, a = unpack_projected_states() + z.copy_(z_generic) + self._forward_core( + mixed_qkv=mixed_qkv, + b=b, + a=a, + core_attn_out=core_attn_out, + ) + return + if gdn_native_fallback in ("1", "true", "all"): + mixed_qkv, z_generic, b, a = unpack_projected_states() + z.copy_(z_generic) + self._forward_core( + mixed_qkv=mixed_qkv, + b=b, + a=a, + core_attn_out=core_attn_out, + ) + return + + if attn_metadata.spec_sequence_masks is not None: # type: ignore[attr-defined] + mixed_qkv, z_generic, b, a = unpack_projected_states() + z.copy_(z_generic) + self._forward_core( + mixed_qkv=mixed_qkv, + b=b, + a=a, + core_attn_out=core_attn_out, + ) + return conv_weights = self.conv1d.weight.view( self.conv1d.weight.size(0), self.conv1d.weight.size(2) ) - torch.ops._xpu_C.gdn_attention( - core_attn_out, - z, - projected_states_qkvz, - projected_states_ba, - self.num_k_heads, - self.num_v_heads, - self.head_k_dim, - self.head_v_dim, - conv_state=self.kv_cache[0], - ssm_state=self.kv_cache[1], - conv_weights=conv_weights, - conv_bias=self.conv1d.bias, - activation=self.activation, - A_log=self.A_log, - dt_bias=self.dt_bias, - num_prefills=attn_metadata.num_prefills, # type: ignore[attr-defined] - num_decodes=attn_metadata.num_decodes, # type: ignore[attr-defined] - has_initial_state=attn_metadata.has_initial_state, # type: ignore[attr-defined] - non_spec_query_start_loc=attn_metadata.non_spec_query_start_loc, # type: ignore[attr-defined] - non_spec_state_indices_tensor=attn_metadata.non_spec_state_indices_tensor, # type: ignore[attr-defined] - num_actual_tokens=attn_metadata.num_actual_tokens, # type: ignore[attr-defined] - tp_size=self.tp_size, - reorder_input=not self.gqa_interleaved_layout, + non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # type: ignore[attr-defined] + if non_spec_state_indices_tensor is not None: + non_spec_state_indices_tensor = non_spec_state_indices_tensor.contiguous() + + non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc # type: ignore[attr-defined] + if non_spec_query_start_loc is not None: + non_spec_query_start_loc = non_spec_query_start_loc.contiguous() + + conv_state = self.kv_cache[0] + ssm_state = self.kv_cache[1] + trace_enabled, trace_context = _xpu_gdn_trace_enabled( + forward_context, self.prefix, attn_metadata ) + if trace_enabled: + _xpu_gdn_trace_write( + stage="pre_native", + trace_context=trace_context, + layer_name=self.prefix, + attn_metadata=attn_metadata, + core_attn_out=core_attn_out, + z=z, + projected_states_qkvz=projected_states_qkvz, + projected_states_ba=projected_states_ba, + conv_state=conv_state, + ssm_state=ssm_state, + non_spec_state_indices_tensor=non_spec_state_indices_tensor, + ) + + if ( + os.environ.get("VLLM_XPU_ZERO_FRESH_GDN_STATE", "0") == "1" + and getattr(attn_metadata, "num_prefills", 0) > 0 + and getattr(attn_metadata, "has_initial_state", None) is not None + and non_spec_state_indices_tensor is not None + ): + try: + has_initial_state = attn_metadata.has_initial_state # type: ignore[attr-defined] + fresh_state_mask = ~has_initial_state + fresh_state_indices = non_spec_state_indices_tensor[fresh_state_mask] + conv_state[fresh_state_indices] = 0 + ssm_state[fresh_state_indices] = 0 + if trace_enabled: + _xpu_gdn_trace_write( + stage="post_fresh_state_zero", + trace_context=trace_context, + layer_name=self.prefix, + attn_metadata=attn_metadata, + core_attn_out=core_attn_out, + z=z, + projected_states_qkvz=projected_states_qkvz, + projected_states_ba=projected_states_ba, + conv_state=conv_state, + ssm_state=ssm_state, + non_spec_state_indices_tensor=non_spec_state_indices_tensor, + ) + except Exception: + logger.exception("Failed to zero fresh XPU GDN state") + + with timed_region("gdn_attention_core_xpu.native"): + torch.ops._xpu_C.gdn_attention( + core_attn_out, + z, + projected_states_qkvz, + projected_states_ba, + self.num_k_heads, + self.num_v_heads, + self.head_k_dim, + self.head_v_dim, + conv_state=conv_state, + ssm_state=ssm_state, + conv_weights=conv_weights, + conv_bias=self.conv1d.bias, + activation=self.activation, + A_log=self.A_log, + dt_bias=self.dt_bias, + num_prefills=attn_metadata.num_prefills, # type: ignore[attr-defined] + num_decodes=attn_metadata.num_decodes, # type: ignore[attr-defined] + has_initial_state=attn_metadata.has_initial_state, # type: ignore[attr-defined] + non_spec_query_start_loc=non_spec_query_start_loc, + non_spec_state_indices_tensor=non_spec_state_indices_tensor, + num_actual_tokens=attn_metadata.num_actual_tokens, # type: ignore[attr-defined] + tp_size=self.tp_size, + reorder_input=not self.gqa_interleaved_layout, + ) + if trace_enabled: + _xpu_gdn_trace_write( + stage="post_native", + trace_context=trace_context, + layer_name=self.prefix, + attn_metadata=attn_metadata, + core_attn_out=core_attn_out, + z=z, + projected_states_qkvz=projected_states_qkvz, + projected_states_ba=projected_states_ba, + conv_state=conv_state, + ssm_state=ssm_state, + non_spec_state_indices_tensor=non_spec_state_indices_tensor, + ) def _gdn_attention_core_xpu_fake( diff --git a/vllm/compilation/cuda_graph.py b/vllm/compilation/cuda_graph.py index b63d86199..1bb32dc62 100644 --- a/vllm/compilation/cuda_graph.py +++ b/vllm/compilation/cuda_graph.py @@ -2,6 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import dataclasses +import json +import os +import re +import time import weakref from collections import Counter from collections.abc import Callable @@ -13,7 +17,10 @@ import torch import vllm.envs as envs from vllm.compilation.counter import compilation_counter -from vllm.compilation.monitor import validate_cudagraph_capturing_enabled +from vllm.compilation.monitor import ( + set_cudagraph_capturing_enabled, + validate_cudagraph_capturing_enabled, +) from vllm.config import CUDAGraphMode, VllmConfig from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id from vllm.forward_context import ( @@ -27,6 +34,34 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import current_stream, weak_ref_tensors logger = init_logger(__name__) +_XPU_CUDAGRAPH_TRACE_LINES = 0 +_XPU_CUDAGRAPH_WRAPPER_IDS = 0 + + +def _xpu_cudagraph_int_env(name: str, default: int) -> int: + try: + return int(os.environ.get(name, str(default)) or default) + except ValueError: + return default + + +def _xpu_cudagraph_rank() -> int: + try: + from vllm.distributed import get_tensor_model_parallel_rank + + return int(get_tensor_model_parallel_rank()) + except Exception: + try: + return int(os.environ.get("LOCAL_RANK", "0") or "0") + except ValueError: + return 0 + + +def _xpu_cudagraph_trace_path(raw_path: str, rank: int) -> str: + try: + return raw_path.format(rank=rank, pid=os.getpid()) + except Exception: + return raw_path @dataclasses.dataclass(frozen=True) @@ -129,6 +164,7 @@ class CUDAGraphEntry: batch_descriptor: BatchDescriptor cudagraph: torch.cuda.CUDAGraph | None = None output: Any | None = None + replay_count: int = 0 # for cudagraph debugging, track the input addresses # during capture, and check if they are the same during replay @@ -182,10 +218,13 @@ class CUDAGraphWrapper: runtime_mode: CUDAGraphMode, cudagraph_options: CUDAGraphOptions | None = None, ) -> None: + global _XPU_CUDAGRAPH_WRAPPER_IDS self.runnable = runnable self.vllm_config = vllm_config self.runtime_mode = runtime_mode self.compilation_config = vllm_config.compilation_config + self._xpu_cudagraph_wrapper_id = _XPU_CUDAGRAPH_WRAPPER_IDS + _XPU_CUDAGRAPH_WRAPPER_IDS += 1 self.first_run_finished = False self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG" @@ -197,10 +236,19 @@ class CUDAGraphWrapper: # TODO: in the future, if we want to use multiple # streams, it might not be safe to share a global pool. # only investigate this when we use multiple streams - self.graph_pool = current_platform.get_global_graph_pool() + if os.environ.get("VLLM_XPU_CUDAGRAPH_PER_WRAPPER_POOL", "0") == "1": + self.graph_pool = current_platform.graph_pool_handle() + elif os.environ.get("VLLM_XPU_CUDAGRAPH_NO_GLOBAL_POOL", "0") == "1": + self.graph_pool = None + else: + self.graph_pool = current_platform.get_global_graph_pool() if cudagraph_options is None: cudagraph_options = CUDAGraphOptions() + if os.environ.get("VLLM_XPU_CUDAGRAPH_STRONG_OUTPUT", "0") == "1": + cudagraph_options = dataclasses.replace( + cudagraph_options, weak_ref_output=False + ) self.cudagraph_options = cudagraph_options # the entries for different batch descriptors that we need to capture # cudagraphs for. @@ -208,6 +256,264 @@ class CUDAGraphWrapper: CUDAGraphWrapper._all_instances.add(self) + def _xpu_cudagraph_identity(self) -> dict[str, Any]: + runnable_type = type(self.runnable) + submod_name = getattr(self.runnable, "submod_name", "") + piecewise_index = getattr(self.runnable, "piecewise_compile_index", None) + total_piecewise = getattr(self.runnable, "total_piecewise_compiles", None) + identity = { + "wrapper_id": self._xpu_cudagraph_wrapper_id, + "runtime_mode": self.runtime_mode.name, + "runnable_type": ( + f"{runnable_type.__module__}.{runnable_type.__qualname__}" + ), + "submod_name": submod_name, + "piecewise_index": piecewise_index, + "total_piecewise_compiles": total_piecewise, + "is_first_graph": getattr(self.runnable, "is_first_graph", None), + "is_last_graph": getattr(self.runnable, "is_last_graph", None), + } + return identity + + def _xpu_cudagraph_trace_label(self) -> str: + identity = self._xpu_cudagraph_identity() + fields = [ + f"wrapper:{identity['wrapper_id']}", + f"mode:{identity['runtime_mode']}", + f"type:{identity['runnable_type']}", + ] + if identity.get("submod_name"): + fields.append(f"submod:{identity['submod_name']}") + if identity.get("piecewise_index") is not None: + fields.append( + "piecewise:" + f"{identity['piecewise_index']}/" + f"{identity.get('total_piecewise_compiles')}" + ) + return " ".join(fields) + + def _xpu_cudagraph_disabled_by_regex(self) -> bool: + pattern = os.environ.get("VLLM_XPU_CUDAGRAPH_DISABLE_SUBMOD_REGEX", "") + if not pattern: + return False + try: + return re.search(pattern, self._xpu_cudagraph_trace_label()) is not None + except re.error as exc: + logger.warning_once( + "Invalid VLLM_XPU_CUDAGRAPH_DISABLE_SUBMOD_REGEX=%r: %r", + pattern, + exc, + ) + return False + + def _xpu_cudagraph_trace_enabled( + self, + forward_context: Any, + ) -> tuple[bool, dict[str, Any]]: + trace_file = os.environ.get("VLLM_XPU_CUDAGRAPH_REPLAY_TRACE_FILE", "") + if not trace_file: + return False, {} + + global _XPU_CUDAGRAPH_TRACE_LINES + max_lines = _xpu_cudagraph_int_env( + "VLLM_XPU_CUDAGRAPH_REPLAY_TRACE_MAX_LINES", 0 + ) + if max_lines > 0 and _XPU_CUDAGRAPH_TRACE_LINES >= max_lines: + return False, {} + + rank = _xpu_cudagraph_rank() + rank_filter = os.environ.get("VLLM_XPU_CUDAGRAPH_REPLAY_TRACE_RANK", "") + if rank_filter and rank_filter not in ("*", "all") and str(rank) != rank_filter: + return False, {} + + label = self._xpu_cudagraph_trace_label() + submod_regex = os.environ.get( + "VLLM_XPU_CUDAGRAPH_REPLAY_TRACE_SUBMOD_REGEX", "" + ) + if submod_regex: + try: + if re.search(submod_regex, label) is None: + return False, {} + except re.error: + return False, {} + + extra = getattr(forward_context, "additional_kwargs", {}) or {} + req_ids = [str(req_id) for req_id in extra.get("xpu_req_ids", [])] + req_regex = os.environ.get("VLLM_XPU_CUDAGRAPH_REPLAY_TRACE_REQ_REGEX", "") + matched_req_ids = req_ids + if req_regex: + try: + matched_req_ids = [ + req_id for req_id in req_ids if re.search(req_regex, req_id) + ] + except re.error: + matched_req_ids = [] + if not matched_req_ids: + return False, {} + + return True, { + "trace_file": _xpu_cudagraph_trace_path(trace_file, rank), + "rank": rank, + "label": label, + "req_ids": req_ids, + "matched_req_ids": matched_req_ids, + "forward_extra": extra, + } + + def _xpu_cudagraph_trace_write( + self, + event: str, + trace_context: dict[str, Any], + batch_descriptor: BatchDescriptor | None, + entry: CUDAGraphEntry | None = None, + args: tuple[Any, ...] | None = None, + output: Any | None = None, + reason: str = "", + ) -> None: + global _XPU_CUDAGRAPH_TRACE_LINES + max_lines = _xpu_cudagraph_int_env( + "VLLM_XPU_CUDAGRAPH_REPLAY_TRACE_MAX_LINES", 0 + ) + if max_lines > 0 and _XPU_CUDAGRAPH_TRACE_LINES >= max_lines: + return + try: + include_args = ( + os.environ.get("VLLM_XPU_CUDAGRAPH_REPLAY_TRACE_INPUTS", "0") + == "1" + ) + include_digest = ( + os.environ.get("VLLM_XPU_CUDAGRAPH_REPLAY_TRACE_DIGEST", "0") + == "1" + ) + record: dict[str, Any] = { + "ts": time.time(), + "pid": os.getpid(), + "tp_rank": trace_context.get("rank"), + "event": event, + "reason": reason, + **self._xpu_cudagraph_identity(), + "label": trace_context.get("label"), + "batch_descriptor": None + if batch_descriptor is None + else str(batch_descriptor), + "req_ids": trace_context.get("req_ids", []), + "matched_req_ids": trace_context.get("matched_req_ids", []), + "forward_extra": trace_context.get("forward_extra", {}), + "entry": None + if entry is None + else { + "has_graph": entry.cudagraph is not None, + "replay_count": int(entry.replay_count), + "input_addresses": entry.input_addresses + if include_args + else None, + }, + } + if include_args and args is not None: + record["tensor_args"] = [ + { + "arg_index": idx, + **self._xpu_cudagraph_tensor_trace_record( + arg, + include_digest=include_digest, + ), + } + for idx, arg in enumerate(args) + if isinstance(arg, torch.Tensor) + ] + if include_digest and output is not None: + record["output"] = self._xpu_cudagraph_output_trace_record(output) + with open(trace_context["trace_file"], "a", encoding="utf-8") as f: + f.write(json.dumps(record, default=str) + "\n") + _XPU_CUDAGRAPH_TRACE_LINES += 1 + except Exception as exc: + logger.warning_once("Failed to write XPU cudagraph replay trace: %r", exc) + + def _xpu_cudagraph_tensor_trace_record( + self, + tensor: torch.Tensor, + *, + include_digest: bool, + ) -> dict[str, Any]: + record: dict[str, Any] = { + "shape": list(tensor.shape), + "dtype": str(tensor.dtype), + "device": str(tensor.device), + "data_ptr": int(tensor.data_ptr()), + "numel": int(tensor.numel()), + } + if not include_digest: + return record + max_numel = _xpu_cudagraph_int_env( + "VLLM_XPU_CUDAGRAPH_REPLAY_TRACE_DIGEST_MAX_NUMEL", 262144 + ) + head_limit = _xpu_cudagraph_int_env( + "VLLM_XPU_CUDAGRAPH_REPLAY_TRACE_DIGEST_HEAD", 8 + ) + if tensor.numel() == 0: + record["digest"] = {"empty": True} + return record + if tensor.numel() > max_numel: + record["digest"] = {"skipped": "numel", "max_numel": max_numel} + return record + try: + with torch.no_grad(): + t = tensor.detach() + stats_t = t.float() if not t.is_floating_point() else t + finite_mask = torch.isfinite(stats_t) + finite_count = int(finite_mask.sum().item()) + digest: dict[str, Any] = { + "finite": finite_count, + "nan": int(torch.isnan(stats_t).sum().item()), + "posinf": int((stats_t == float("inf")).sum().item()), + "neginf": int((stats_t == -float("inf")).sum().item()), + } + if finite_count: + finite_vals = stats_t[finite_mask].float() + digest.update( + { + "min": float(finite_vals.min().item()), + "max": float(finite_vals.max().item()), + "mean": float(finite_vals.mean().item()), + "sum": float(finite_vals.sum().item()), + "l2": float(torch.linalg.vector_norm(finite_vals).item()), + } + ) + if head_limit > 0: + digest["head"] = ( + t.reshape(-1)[:head_limit].to("cpu").tolist() + ) + record["digest"] = digest + except Exception as exc: + record["digest"] = {"error": repr(exc)} + return record + + def _xpu_cudagraph_output_trace_record(self, output: Any) -> Any: + max_tensors = _xpu_cudagraph_int_env( + "VLLM_XPU_CUDAGRAPH_REPLAY_TRACE_OUTPUT_MAX_TENSORS", 8 + ) + tensors_seen = 0 + + def visit(obj: Any) -> Any: + nonlocal tensors_seen + if isinstance(obj, torch.Tensor): + if tensors_seen >= max_tensors: + return {"skipped": "max_tensors"} + tensors_seen += 1 + return self._xpu_cudagraph_tensor_trace_record( + obj, + include_digest=True, + ) + if isinstance(obj, tuple): + return [visit(item) for item in obj] + if isinstance(obj, list): + return [visit(item) for item in obj] + if isinstance(obj, dict): + return {str(key): visit(value) for key, value in obj.items()} + return str(type(obj)) + + return visit(output) + def __getattr__(self, key: str) -> Any: # allow accessing the attributes of the runnable. if hasattr(self.runnable, key): @@ -230,6 +536,14 @@ class CUDAGraphWrapper: def clear_graphs(self) -> None: self.concrete_cudagraph_entries.clear() + @staticmethod + def _recapture_after_n_replays() -> int: + raw = os.environ.get("VLLM_XPU_CUDAGRAPH_RECAPTURE_AFTER_N_REPLAYS", "0") + try: + return max(0, int(raw)) + except ValueError: + return 0 + def __call__(self, *args: Any, **kwargs: Any) -> Any | None: if not is_forward_context_available(): # No forward context means we are outside the normal @@ -240,6 +554,9 @@ class CUDAGraphWrapper: forward_context = get_forward_context() batch_descriptor = forward_context.batch_descriptor cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode + trace_enabled, trace_context = self._xpu_cudagraph_trace_enabled( + forward_context + ) if ( cudagraph_runtime_mode == CUDAGraphMode.NONE @@ -251,9 +568,54 @@ class CUDAGraphWrapper: # matches. This enables properly dispatching to the correct # CUDAGraphWrapper when nesting multiple instances with different # runtime modes. - return self.runnable(*args, **kwargs) + if trace_enabled: + self._xpu_cudagraph_trace_write( + "direct_start", + trace_context, + batch_descriptor, + args=args, + reason=( + f"runtime_mode={cudagraph_runtime_mode.name} " + f"wrapper_mode={self.runtime_mode.name}" + ), + ) + output = self.runnable(*args, **kwargs) + if trace_enabled: + self._xpu_cudagraph_trace_write( + "direct_finish", + trace_context, + batch_descriptor, + args=args, + output=output, + reason=( + f"runtime_mode={cudagraph_runtime_mode.name} " + f"wrapper_mode={self.runtime_mode.name}" + ), + ) + return output assert batch_descriptor is not None + if self._xpu_cudagraph_disabled_by_regex(): + if trace_enabled: + self._xpu_cudagraph_trace_write( + "direct_start", + trace_context, + batch_descriptor, + args=args, + reason="disabled_by_submod_regex", + ) + output = self.runnable(*args, **kwargs) + if trace_enabled: + self._xpu_cudagraph_trace_write( + "direct_finish", + trace_context, + batch_descriptor, + args=args, + output=output, + reason="disabled_by_submod_regex", + ) + return output + if batch_descriptor not in self.concrete_cudagraph_entries: # create a new entry for this batch descriptor self.concrete_cudagraph_entries[batch_descriptor] = CUDAGraphEntry( @@ -281,6 +643,14 @@ class CUDAGraphWrapper: ] entry.input_addresses = input_addresses cudagraph = torch.cuda.CUDAGraph() + if trace_enabled: + self._xpu_cudagraph_trace_write( + "capture_start", + trace_context, + batch_descriptor, + entry, + args=args, + ) with ExitStack() as stack: if self.cudagraph_options.gc_disable: @@ -337,6 +707,15 @@ class CUDAGraphWrapper: entry.cudagraph = cudagraph compilation_counter.num_cudagraph_captured += 1 + if trace_enabled: + self._xpu_cudagraph_trace_write( + "capture_finish", + trace_context, + batch_descriptor, + entry, + args=args, + output=output, + ) # important: we need to return the output, rather than # the weak ref of the output, so that pytorch can correctly @@ -357,5 +736,55 @@ class CUDAGraphWrapper: # Sync offloader before replay - ensures any external dependencies # from pre-capture prefetches are satisfied. get_offloader().sync_prev_onload() + recapture_after = self._recapture_after_n_replays() + if recapture_after and entry.replay_count >= recapture_after: + if trace_enabled: + self._xpu_cudagraph_trace_write( + "recapture_requested", + trace_context, + batch_descriptor, + entry, + args=args, + reason=f"replay_count>={recapture_after}", + ) + entry.cudagraph = None + entry.output = None + entry.input_addresses = None + entry.replay_count = 0 + if ( + os.environ.get( + "VLLM_XPU_CUDAGRAPH_ALLOW_RUNTIME_RECAPTURE", "0" + ) + == "1" + ): + set_cudagraph_capturing_enabled(True) + try: + return self.__call__(*args, **kwargs) + finally: + set_cudagraph_capturing_enabled(False) + return self.__call__(*args, **kwargs) + sync_replay = os.environ.get("VLLM_XPU_SYNC_CUDAGRAPH_REPLAY", "0") == "1" + if sync_replay and hasattr(torch, "xpu"): + torch.xpu.synchronize() + if trace_enabled: + self._xpu_cudagraph_trace_write( + "replay_start", + trace_context, + batch_descriptor, + entry, + args=args, + ) entry.cudagraph.replay() + entry.replay_count += 1 + if sync_replay and hasattr(torch, "xpu"): + torch.xpu.synchronize() + if trace_enabled: + self._xpu_cudagraph_trace_write( + "replay_finish", + trace_context, + batch_descriptor, + entry, + args=args, + output=entry.output, + ) return entry.output diff --git a/vllm/model_executor/kernels/linear/scaled_mm/xpu.py b/vllm/model_executor/kernels/linear/scaled_mm/xpu.py index 6d75a420e..32f989a53 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/xpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/xpu.py @@ -2,14 +2,29 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Sequence +import os import torch +import torch.nn.functional as F +from vllm import _custom_ops as ops from vllm.model_executor.kernels.linear import ( # noqa: E501 FP8ScaledMMLinearKernel, FP8ScaledMMLinearLayerConfig, + Int8ScaledMMLinearKernel, + Int8ScaledMMLinearLayerConfig, +) +from vllm.model_executor.kernels.linear.scaled_mm.BlockScaledMMLinearKernel import ( + Fp8BlockScaledMMLinearKernel, +) +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + process_fp8_weight_block_strategy, +) +from vllm.model_executor.layers.quantization.utils.int8_utils import ( + per_token_quant_int8, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8Static128BlockSym, kFp8StaticChannelSym, kFp8StaticTensorSym, ) @@ -17,6 +32,180 @@ from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform +_INT8_GEMM_W8A8_FAKE_REGISTERED = False +_INT8_PER_TOKEN_QUANT_FAKE_REGISTERED = False + + +def _register_int8_gemm_w8a8_fake() -> None: + global _INT8_GEMM_W8A8_FAKE_REGISTERED + if _INT8_GEMM_W8A8_FAKE_REGISTERED: + return + + def _fake_int8_gemm_w8a8( + A: torch.Tensor, + A_scale: torch.Tensor, + B: torch.Tensor, + B_scale: torch.Tensor, + out_dtype: torch.dtype | None, + bias: torch.Tensor | None, + ) -> torch.Tensor: + del A_scale, B_scale, bias + return torch.empty( + (A.shape[0], B.shape[1]), + device=A.device, + dtype=out_dtype or torch.bfloat16, + ) + + try: + torch.library.register_fake("_xpu_C::int8_gemm_w8a8")( + _fake_int8_gemm_w8a8 + ) + except RuntimeError as e: + if "already has" not in str(e) and "already registered" not in str(e): + raise + _INT8_GEMM_W8A8_FAKE_REGISTERED = True + + +def _register_int8_per_token_quant_fake() -> None: + global _INT8_PER_TOKEN_QUANT_FAKE_REGISTERED + if _INT8_PER_TOKEN_QUANT_FAKE_REGISTERED: + return + + def _fake_per_token_quant_int8_xpu( + x: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + q = torch.empty_like(x, dtype=torch.int8) + scales = torch.empty( + (*x.shape[:-1], 1), + device=x.device, + dtype=torch.float32, + ) + return q, scales + + try: + torch.library.register_fake("_xpu_C::per_token_quant_int8_xpu")( + _fake_per_token_quant_int8_xpu + ) + except RuntimeError as e: + if "already has" not in str(e) and "already registered" not in str(e): + raise + _INT8_PER_TOKEN_QUANT_FAKE_REGISTERED = True + + +class XPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel): + """XPU W8A8 INT8 dense GEMM for dynamic symmetric per-token activations.""" + + @staticmethod + def _has_int8_gemm_w8a8() -> bool: + try: + import vllm_xpu_kernels._xpu_C # noqa: F401 + except Exception: + return False + try: + torch._C._dispatch_find_schema_or_throw( + "_xpu_C::int8_gemm_w8a8", + "", + ) + _register_int8_gemm_w8a8_fake() + return True + except RuntimeError: + return False + + @staticmethod + def _has_per_token_quant_int8_xpu() -> bool: + try: + import vllm_xpu_kernels._xpu_C # noqa: F401 + except Exception: + return False + try: + torch._C._dispatch_find_schema_or_throw( + "_xpu_C::per_token_quant_int8_xpu", + "", + ) + _register_int8_per_token_quant_fake() + return True + except RuntimeError: + return False + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_xpu(): + return False, "XPUInt8ScaledMM only supports XPU" + if not cls._has_int8_gemm_w8a8(): + return False, "vllm-xpu-kernels does not expose int8_gemm_w8a8" + return True, None + + @classmethod + def can_implement( + cls, c: Int8ScaledMMLinearLayerConfig + ) -> tuple[bool, str | None]: + if c.is_static_input_scheme: + return False, "only dynamic activation scales are supported" + if not c.input_symmetric: + return False, "only symmetric activation quantization is supported" + if not c.is_channelwise: + return False, "only per-channel weight scales are supported" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + w_q_name, w_s_name, i_s_name, i_zp_name, azp_adj_name = self.layer_param_names + + weight = getattr(layer, w_q_name) + if os.getenv("VLLM_XPU_INT8_LINEAR_BF16_FALLBACK", "0") == "1": + weight_scale = getattr(layer, w_s_name) + scale = weight_scale.data.to(torch.bfloat16) + while scale.ndim < weight.ndim: + scale = scale.unsqueeze(-1) + weight_bf16 = weight.data.to(torch.bfloat16) * scale + replace_parameter(layer, w_q_name, weight_bf16.contiguous()) + setattr(layer, i_s_name, None) + setattr(layer, i_zp_name, None) + setattr(layer, azp_adj_name, None) + layer.xpu_int8_linear_bf16_fallback = True + layer.xpu_native_int8_activation_quant = False + return + + replace_parameter(layer, w_q_name, weight.t().contiguous()) + + weight_scale = getattr(layer, w_s_name) + replace_parameter(layer, w_s_name, weight_scale.flatten().contiguous()) + + setattr(layer, i_s_name, None) + setattr(layer, i_zp_name, None) + setattr(layer, azp_adj_name, None) + layer.xpu_native_int8_activation_quant = ( + self._has_per_token_quant_int8_xpu() + and os.getenv("VLLM_XPU_DISABLE_NATIVE_INT8_ACTIVATION_QUANT", "0") + != "1" + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + w_q, w_s, _, _, _ = self._get_layer_params(layer) + if getattr(layer, "xpu_int8_linear_bf16_fallback", False): + return F.linear(x, w_q, bias) + x_2d = x.view(-1, x.shape[-1]).contiguous() + if getattr(layer, "xpu_native_int8_activation_quant", False): + x_q, x_s = torch.ops._xpu_C.per_token_quant_int8_xpu(x_2d) + else: + x_q, x_s = per_token_quant_int8(x_2d) + out = torch.ops._xpu_C.int8_gemm_w8a8( + x_q, + x_s, + w_q, + w_s, + x.dtype, + bias, + ) + return out.view(*x.shape[:-1], w_q.shape[1]) + + class XPUFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): @classmethod def is_supported( @@ -46,6 +235,15 @@ class XPUFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): self.layer_param_names = layer_param_names def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if os.getenv("VLLM_XPU_FP8_LINEAR_BF16_FALLBACK", "0") == "1": + weight = layer.weight.data + weight_scale = layer.weight_scale.data.to(torch.bfloat16) + while weight_scale.ndim < weight.ndim: + weight_scale = weight_scale.unsqueeze(-1) + weight_bf16 = weight.to(torch.bfloat16) * weight_scale + replace_parameter(layer, "weight", weight_bf16.contiguous()) + layer.input_scale = None + return replace_parameter(layer, "weight", layer.weight.data.t()) def apply_weights( @@ -54,6 +252,8 @@ class XPUFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: + if os.getenv("VLLM_XPU_FP8_LINEAR_BF16_FALLBACK", "0") == "1": + return F.linear(x, layer.weight, bias) weight = layer.weight weight_scale = layer.weight_scale return torch.ops._xpu_C.fp8_gemm_w8a16(x, weight, weight_scale, bias) @@ -70,3 +270,174 @@ class XPUFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): output_shape: list, ) -> torch.Tensor: pass + + +class XPUBF16Fp8BlockScaledMMLinearKernel(Fp8BlockScaledMMLinearKernel): + """BF16 fallback for block-FP8 checkpoints on XPU. + + Intel's current XPU kernel set exposes FP8 weight-only GEMM, but not the + 128x128 block-scaled W8A8 GEMM used by Qwen3.6 FP8 checkpoints. This path + keeps the checkpoint's FP8 values and scales exact by dequantizing weights + once after load, then using the regular BF16 matmul path. + """ + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_xpu(): + return False, "XPUBF16Fp8BlockScaledMM only supports XPU" + return True, None + + @classmethod + def can_implement( + cls, config: FP8ScaledMMLinearLayerConfig + ) -> tuple[bool, str | None]: + can_implement, failure_reason = super().can_implement(config) + if not can_implement: + return can_implement, failure_reason + + if config.weight_quant_key != kFp8Static128BlockSym: + return ( + False, + "XPUBF16Fp8BlockScaledMM only supports static 128x128 FP8 blocks", + ) + if config.weight_quant_key.dtype not in { + torch.float8_e5m2, + torch.float8_e4m3fn, + }: + return False, "XPUBF16Fp8BlockScaledMM only supports FP8 weight dtype" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + params = self._get_layer_params(layer) + weight_scale = ( + params.weight_scale + if params.weight_scale_inv is None + else params.weight_scale_inv + ) + assert weight_scale is not None + + weight, weight_scale = process_fp8_weight_block_strategy( + params.weight, + weight_scale, + ) + block_n, block_k = self.weight_group_shape + scale = weight_scale.to(torch.bfloat16) + expanded_scale = scale.repeat_interleave(block_n, dim=0).repeat_interleave( + block_k, dim=1 + ) + expanded_scale = expanded_scale[: weight.shape[0], : weight.shape[1]] + weight_bf16 = weight.to(torch.bfloat16) * expanded_scale + + replace_parameter(layer, params.WEIGHT, weight_bf16.contiguous()) + layer.input_scale = None + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + return F.linear(x, layer.weight, bias) + + def apply_block_scaled_mm( + self, + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + ) -> torch.Tensor: + raise NotImplementedError("XPUBF16Fp8BlockScaledMM uses apply_weights") + + +class XPURequantFp8BlockScaledMMLinearKernel(Fp8BlockScaledMMLinearKernel): + """Experimental XPU path that maps block-FP8 checkpoints to W8A16 FP8 GEMM. + + This is faster and more memory efficient than the BF16 fallback, but it + requantizes the checkpoint from 128x128 block-scaled FP8 into per-channel + FP8 weights. Keep it opt-in until quality impact is measured. + """ + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_xpu(): + return False, "XPURequantFp8BlockScaledMM only supports XPU" + if os.getenv("VLLM_XPU_BLOCK_FP8_REQUANT", "0") != "1": + return False, "set VLLM_XPU_BLOCK_FP8_REQUANT=1 to enable" + return True, None + + @classmethod + def can_implement( + cls, config: FP8ScaledMMLinearLayerConfig + ) -> tuple[bool, str | None]: + can_implement, failure_reason = super().can_implement(config) + if not can_implement: + return can_implement, failure_reason + + if config.weight_quant_key != kFp8Static128BlockSym: + return ( + False, + "XPURequantFp8BlockScaledMM only supports static 128x128 FP8 blocks", + ) + if config.weight_quant_key.dtype not in { + torch.float8_e5m2, + torch.float8_e4m3fn, + }: + return False, "XPURequantFp8BlockScaledMM only supports FP8 weight dtype" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + params = self._get_layer_params(layer) + weight_scale = ( + params.weight_scale + if params.weight_scale_inv is None + else params.weight_scale_inv + ) + assert weight_scale is not None + + weight, weight_scale = process_fp8_weight_block_strategy( + params.weight, + weight_scale, + ) + block_n, block_k = self.weight_group_shape + scale = weight_scale.to(torch.bfloat16) + expanded_scale = scale.repeat_interleave(block_n, dim=0).repeat_interleave( + block_k, dim=1 + ) + expanded_scale = expanded_scale[: weight.shape[0], : weight.shape[1]] + weight_bf16 = weight.to(torch.bfloat16) * expanded_scale + + qweight, qscale = ops.scaled_fp8_quant( + weight_bf16.contiguous(), + use_per_token_if_dynamic=True, + ) + replace_parameter(layer, params.WEIGHT, qweight.t().contiguous()) + replace_parameter(layer, params.WEIGHT_SCALE_INV, qscale.flatten().contiguous()) + layer.input_scale = None + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + return torch.ops._xpu_C.fp8_gemm_w8a16( + x, + layer.weight, + layer.weight_scale_inv, + bias, + ) + + def apply_block_scaled_mm( + self, + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + ) -> torch.Tensor: + raise NotImplementedError("XPURequantFp8BlockScaledMM uses apply_weights") diff --git a/vllm/model_executor/layers/mamba/gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn_linear_attn.py index a621ab962..d7c637fbd 100644 --- a/vllm/model_executor/layers/mamba/gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn_linear_attn.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Inference-only Qwen3-Next/Qwen3.5 model.""" +import os import torch from einops import rearrange from torch import nn @@ -62,6 +63,7 @@ from vllm.utils.torch_utils import ( _resolve_layer_name, direct_register_custom_op, ) +from vllm.utils.xpu_decode_timing import timed_region from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata logger = init_logger(__name__) @@ -384,6 +386,14 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): self.enable_packed_recurrent_decode = ( envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE ) + xpu_gdn_reuse_qkvz_ba_quant_mode = ( + envs.VLLM_XPU_GDN_REUSE_QKVZ_BA_QUANT + ) + self.xpu_reuse_qkvz_ba_quant_mode = xpu_gdn_reuse_qkvz_ba_quant_mode + self.xpu_reuse_qkvz_ba_quant = current_platform.is_xpu() and ( + xpu_gdn_reuse_qkvz_ba_quant_mode + in ("1", "true", "clone", "clone-ba", "clone-qkvz") + ) compilation_config = get_current_vllm_config().compilation_config if prefix in compilation_config.static_forward_context: @@ -612,8 +622,61 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): # ============================================================ # Part 1: Input Projection # ============================================================ - projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states) - projected_states_ba, _ = self.in_proj_ba(hidden_states) + if self.xpu_reuse_qkvz_ba_quant and ( + getattr(self.in_proj_qkvz, "xpu_native_int8_activation_quant", False) + and getattr(self.in_proj_ba, "xpu_native_int8_activation_quant", False) + and self.in_proj_qkvz.bias is None + and self.in_proj_ba.bias is None + and not self.in_proj_qkvz.gather_output + and not self.in_proj_ba.gather_output + ): + x_2d = hidden_states.view(-1, hidden_states.shape[-1]).contiguous() + with timed_region("qwen3_next.gdn.input_quant"): + x_q, x_s = torch.ops._xpu_C.per_token_quant_int8_xpu(x_2d) + if self.xpu_reuse_qkvz_ba_quant_mode == "clone": + x_q_qkvz = x_q.clone() + x_s_qkvz = x_s.clone() + x_q_ba = x_q.clone() + x_s_ba = x_s.clone() + elif self.xpu_reuse_qkvz_ba_quant_mode == "clone-ba": + x_q_qkvz = x_q + x_s_qkvz = x_s + x_q_ba = x_q.clone() + x_s_ba = x_s.clone() + elif self.xpu_reuse_qkvz_ba_quant_mode == "clone-qkvz": + x_q_qkvz = x_q.clone() + x_s_qkvz = x_s.clone() + x_q_ba = x_q + x_s_ba = x_s + else: + x_q_qkvz = x_q + x_s_qkvz = x_s + x_q_ba = x_q + x_s_ba = x_s + + with timed_region("qwen3_next.gdn.qkvz_gemm_w8a8"): + projected_states_qkvz = torch.ops._xpu_C.int8_gemm_w8a8( + x_q_qkvz, + x_s_qkvz, + self.in_proj_qkvz.weight, + self.in_proj_qkvz.weight_scale, + hidden_states.dtype, + None, + ).view(*hidden_states.shape[:-1], self.in_proj_qkvz.weight.shape[1]) + with timed_region("qwen3_next.gdn.ba_gemm_w8a8"): + projected_states_ba = torch.ops._xpu_C.int8_gemm_w8a8( + x_q_ba, + x_s_ba, + self.in_proj_ba.weight, + self.in_proj_ba.weight_scale, + hidden_states.dtype, + None, + ).view(*hidden_states.shape[:-1], self.in_proj_ba.weight.shape[1]) + else: + with timed_region("qwen3_next.gdn.qkvz_proj"): + projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states) + with timed_region("qwen3_next.gdn.ba_proj"): + projected_states_ba, _ = self.in_proj_ba(hidden_states) # ============================================================ # Part 2: Core Attention @@ -625,13 +688,14 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): ) z = torch.empty_like(core_attn_out) - torch.ops.vllm.gdn_attention_core_xpu( - core_attn_out, - z, - projected_states_qkvz, - projected_states_ba, - self.prefix, - ) + with timed_region("qwen3_next.gdn.core_op"): + torch.ops.vllm.gdn_attention_core_xpu( + core_attn_out, + z, + projected_states_qkvz, + projected_states_ba, + self.prefix, + ) # ============================================================ # Part 3: Output Projection @@ -640,10 +704,12 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): # Reshape input data into 2D tensor core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1]) z = z.reshape(-1, z.shape[-1]) - core_attn_out = self.norm(core_attn_out, z) + with timed_region("qwen3_next.gdn.output_norm"): + core_attn_out = self.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(z_shape_og) core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)") - output[:num_tokens], _ = self.out_proj(core_attn_out) + with timed_region("qwen3_next.gdn.out_proj"): + output[:num_tokens], _ = self.out_proj(core_attn_out) def _warmup_prefill_kernels(self, mixed_qkv: torch.Tensor) -> None: """Warm up GDN prefill kernels during V1 profiling. @@ -845,6 +911,15 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): # 1.2: Process the remaining part if attn_metadata.num_prefills > 0: assert mixed_qkv_non_spec is not None + if ( + os.environ.get("VLLM_XPU_ZERO_FRESH_GDN_STATE", "0") == "1" + and has_initial_state is not None + and non_spec_state_indices_tensor is not None + ): + fresh_state_mask = ~has_initial_state + fresh_state_indices = non_spec_state_indices_tensor[fresh_state_mask] + conv_state[fresh_state_indices] = 0 + ssm_state[fresh_state_indices] = 0 mixed_qkv_non_spec_T = mixed_qkv_non_spec.transpose(0, 1) # - "cache_indices" updates the conv_state cache in positions # pointed to by "state_indices_tensor" diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 1eeca1423..c06ee9cb0 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -3,7 +3,9 @@ from typing import Any +import os import torch +import torch.nn.functional as F import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm import _custom_ops as ops @@ -37,6 +39,11 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( mxfp4_round_up_hidden_size_and_intermediate_size, select_gpt_oss_mxfp4_moe_backend, ) +from vllm.model_executor.layers.fused_moe.oracle.int8 import ( + Int8MoeBackend, + make_int8_moe_kernel, + select_int8_moe_backend, +) from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( prepare_fp8_moe_layer_for_marlin, ) @@ -47,7 +54,11 @@ from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( OCP_MX_BLOCK_SIZE, OCP_MX_Scheme, ) -from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + kInt8DynamicTokenSym, + kInt8StaticChannelSym, +) from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( all_close_1d, normalize_e4m3fn_to_e4m3fnuz, @@ -59,6 +70,104 @@ from vllm.scalar_type import scalar_types logger = init_logger(__name__) + +def _xpu_moe_diagnostic_context_enabled() -> bool: + return bool( + os.environ.get("VLLM_XPU_MOE_LIVE_ABI_FILE") + or os.environ.get("VLLM_XPU_MOE_ONEDNN_SIDECAR_PROBE") + or os.environ.get("VLLM_XPU_MOE_REPLAY_DIGEST") + ) + + +def _safe_layer_index(layer: FusedMoE) -> int: + try: + return int(layer.layer_id) + except Exception: + return -1 + + +def _xpu_int8_moe_bf16_fallback_mode() -> str: + mode = os.environ.get("VLLM_XPU_INT8_MOE_BF16_FALLBACK", "0") + return mode.strip().lower() + + +def _xpu_int8_moe_bf16_fallback_enabled() -> bool: + return _xpu_int8_moe_bf16_fallback_mode() in ( + "1", + "all", + "decode", + "decode_only", + "true", + "yes", + "on", + ) + + +def _dequant_int8_moe_weight_bf16( + weight: torch.Tensor, + scale: torch.Tensor, +) -> torch.Tensor: + scale_bf16 = scale.to(torch.bfloat16) + while scale_bf16.dim() < weight.dim(): + scale_bf16 = scale_bf16.unsqueeze(-1) + return weight.to(torch.bfloat16) * scale_bf16 + + +def _apply_quark_int8_moe_bf16_fallback( + layer: FusedMoE, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + apply_router_weight_on_input: bool, +) -> torch.Tensor: + w13 = layer.w13_weight_bf16 + w2 = layer.w2_weight_bf16 + topk = topk_ids.shape[-1] + flat_ids = topk_ids.reshape(-1) + + if layer.expert_map is not None: + local_ids = layer.expert_map[flat_ids.to(torch.long)] + else: + local_ids = flat_ids + local_ids = local_ids.to(torch.long) + valid = (local_ids >= 0) & (local_ids < w13.shape[0]) + safe_ids = torch.clamp(local_ids, min=0, max=w13.shape[0] - 1) + + routed = x.to(torch.bfloat16).repeat_interleave(topk, dim=0) + route_weights = topk_weights.reshape(-1, 1).to(torch.bfloat16) + if apply_router_weight_on_input: + routed = routed * route_weights + + hidden = torch.bmm(w13[safe_ids], routed.unsqueeze(-1)).squeeze(-1) + if getattr(layer, "w13_bias", None) is not None: + hidden = hidden + layer.w13_bias[safe_ids].to(torch.bfloat16) + + if layer.activation == MoEActivation.SILU: + half = hidden.shape[-1] // 2 + hidden = F.silu(hidden[..., :half]) * hidden[..., half:] + elif layer.activation == MoEActivation.GELU: + half = hidden.shape[-1] // 2 + hidden = F.gelu(hidden[..., :half]) * hidden[..., half:] + elif layer.activation == MoEActivation.SILU_NO_MUL: + hidden = F.silu(hidden) + elif layer.activation == MoEActivation.GELU_NO_MUL: + hidden = F.gelu(hidden) + elif layer.activation == MoEActivation.RELU2_NO_MUL: + hidden = torch.square(F.relu(hidden)) + else: + raise ValueError( + "VLLM_XPU_INT8_MOE_BF16_FALLBACK does not support " + f"activation {layer.activation.value!r}." + ) + + routed_out = torch.bmm(w2[safe_ids], hidden.unsqueeze(-1)).squeeze(-1) + if getattr(layer, "w2_bias", None) is not None: + routed_out = routed_out + layer.w2_bias[safe_ids].to(torch.bfloat16) + if not apply_router_weight_on_input: + routed_out = routed_out * route_weights + routed_out = routed_out * valid.reshape(-1, 1).to(torch.bfloat16) + return routed_out.reshape(x.shape[0], topk, -1).sum(dim=1).to(x.dtype) + __all__ = [ "QuarkMoEMethod", "QuarkOCP_MX_MoEMethod", @@ -526,6 +635,13 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod): self.input_quant = input_config self.weight_qscheme = self.weight_quant.get("qscheme", "per_tensor") self.static_input_scales = not self.input_quant.get("is_dynamic", False) + activation_key = None if self.static_input_scales else kInt8DynamicTokenSym + self.int8_backend, self.experts_cls = select_int8_moe_backend( + config=self.moe, + weight_key=kInt8StaticChannelSym, + activation_key=activation_key, + ) + self.moe_kernel: mk.FusedMoEKernel | None = None def create_weights( self, @@ -745,6 +861,52 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod): max_w13_scales, requires_grad=False ) + if ( + self.int8_backend == Int8MoeBackend.XPU + and _xpu_int8_moe_bf16_fallback_enabled() + ): + layer.xpu_int8_moe_bf16_fallback = True + layer.w13_weight_bf16 = torch.nn.Parameter( + _dequant_int8_moe_weight_bf16( + layer.w13_weight.data, + layer.w13_weight_scale.data, + ), + requires_grad=False, + ) + layer.w2_weight_bf16 = torch.nn.Parameter( + _dequant_int8_moe_weight_bf16( + layer.w2_weight.data, + layer.w2_weight_scale.data, + ), + requires_grad=False, + ) + + if self.int8_backend == Int8MoeBackend.XPU: + from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( + prepare_int8_moe_layer_for_xpu, + ) + + w13, w2, w13_scale, w2_scale = prepare_int8_moe_layer_for_xpu( + layer.w13_weight, + layer.w2_weight, + layer.w13_weight_scale, + layer.w2_weight_scale, + ) + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w13_weight_scale", w13_scale) + replace_parameter(layer, "w2_weight_scale", w2_scale) + + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + if self.int8_backend == Int8MoeBackend.XPU: + self.moe_kernel = make_int8_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + experts_cls=self.experts_cls, + routing_tables=layer._maybe_init_expert_routing_tables(), + shared_experts=layer.shared_experts, + ) + def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> FusedMoEQuantConfig | None: @@ -771,6 +933,64 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod): topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if getattr(layer, "xpu_int8_moe_bf16_fallback", False): + mode = _xpu_int8_moe_bf16_fallback_mode() + use_fallback = mode not in ("decode", "decode_only") or x.shape[0] == 1 + if use_fallback: + return _apply_quark_int8_moe_bf16_fallback( + layer, + x, + topk_weights, + topk_ids, + layer.apply_router_weight_on_input, + ) + + if self.int8_backend == Int8MoeBackend.XPU: + assert self.moe_kernel is not None + route_capture = getattr(layer, "route_capture", None) + if route_capture is not None: + route_capture.capture(topk_ids, stage="quark_int8_apply") + fused_experts = getattr(self.moe_kernel, "fused_experts", None) + had_live_abi_context = False + old_live_abi_context = None + apply_router_weight_on_input = layer.apply_router_weight_on_input + if (fused_experts is not None + and _xpu_moe_diagnostic_context_enabled()): + had_live_abi_context = hasattr(fused_experts, + "_live_abi_context") + old_live_abi_context = getattr(fused_experts, + "_live_abi_context", None) + layer_index = _safe_layer_index(layer) + fused_experts._live_abi_context = { + "layer": getattr(layer, "layer_name", ""), + "layer_index": layer_index, + "global_num_experts": int(layer.global_num_experts), + "local_num_experts": int(layer.local_num_experts), + "activation": layer.activation.value, + "apply_router_weight_on_input": + bool(apply_router_weight_on_input), + "int8_backend": self.int8_backend.value, + } + try: + return self.moe_kernel.apply( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=apply_router_weight_on_input, + shared_experts_input=shared_experts_input, + ) + finally: + if fused_experts is not None: + if had_live_abi_context: + fused_experts._live_abi_context = old_live_abi_context + elif hasattr(fused_experts, "_live_abi_context"): + delattr(fused_experts, "_live_abi_context") + from vllm.model_executor.layers.fused_moe import fused_experts return fused_experts( diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py index a77eafba2..a3dd44e69 100644 --- a/vllm/v1/sample/sampler.py +++ b/vllm/v1/sample/sampler.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """A layer that samples the next tokens from the model's outputs.""" +import os + import torch import torch.nn as nn @@ -16,6 +18,36 @@ from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler from vllm.v1.worker.gpu.sample.logprob import compute_token_logprobs _SAMPLING_EPS = 1e-5 +_XPU_SAMPLER_DEVICE_TIMING = ( + os.environ.get("VLLM_XPU_SAMPLER_DEVICE_TIMING", "0") == "1" +) + + +def _xpu_sampler_record_event(device_type: str): + if not _XPU_SAMPLER_DEVICE_TIMING: + return None + try: + if device_type == "xpu" and hasattr(torch, "xpu"): + event = torch.xpu.Event(enable_timing=True) + elif device_type == "cuda" and torch.cuda.is_available(): + event = torch.cuda.Event(enable_timing=True) + else: + return None + event.record() + return event + except Exception: + return None + + +def _xpu_sampler_mark(timeline: dict[str, object] | None, key: str) -> None: + if timeline is None: + return + device_type = timeline.get("device_type") + if not isinstance(device_type, str): + return + event = _xpu_sampler_record_event(device_type) + if event is not None: + timeline[key] = event class Sampler(nn.Module): @@ -73,6 +105,28 @@ class Sampler(nn.Module): logprobs_mode_override: LogprobsMode | None = None, ) -> SamplerOutput: logprobs_mode = logprobs_mode_override or self.logprobs_mode + sampler_timeline: dict[str, object] | None = None + if _XPU_SAMPLER_DEVICE_TIMING: + device_type = getattr(logits.device, "type", "") + if device_type in ("xpu", "cuda"): + sampler_timeline = { + "device_type": device_type, + "sampler_logprobs_mode": str(logprobs_mode), + "sampler_num_logprobs": sampling_metadata.max_num_logprobs, + "sampler_all_greedy": bool(sampling_metadata.all_greedy), + "sampler_all_random": bool(sampling_metadata.all_random), + "sampler_no_penalties": bool(sampling_metadata.no_penalties), + "sampler_has_allowed_token_mask": ( + sampling_metadata.allowed_token_ids_mask is not None + ), + "sampler_bad_words_count": len( + sampling_metadata.bad_words_token_ids or () + ), + "sampler_logprob_token_ids_count": len( + sampling_metadata.logprob_token_ids or {} + ), + } + _xpu_sampler_mark(sampler_timeline, "sampler_entry_event") # NOTE(woosuk): Use the original logits (before any penalties or # temperature scaling) for the top-k logprobs. # This is different from the V0 sampler, which uses the logits that @@ -86,15 +140,23 @@ class Sampler(nn.Module): raw_logprobs = logits.clone() else: raw_logprobs = logits.to(torch.float32) + _xpu_sampler_mark(sampler_timeline, "sampler_raw_logprobs_end_event") # Use float32 for the logits. logits = logits.to(torch.float32) + _xpu_sampler_mark(sampler_timeline, "sampler_fp32_cast_end_event") logits = self.apply_logits_processors( logits, sampling_metadata, predict_bonus_token ) + _xpu_sampler_mark(sampler_timeline, "sampler_processors_end_event") # Sample the next token. - sampled, processed_logprobs = self.sample(logits, sampling_metadata) + self._xpu_sampler_timeline = sampler_timeline + try: + sampled, processed_logprobs = self.sample(logits, sampling_metadata) + finally: + self._xpu_sampler_timeline = None + _xpu_sampler_mark(sampler_timeline, "sampler_sample_call_end_event") if processed_logprobs is not None: raw_logprobs = processed_logprobs # Convert sampled token ids to int64 (long) type to ensure compatibility @@ -102,6 +164,7 @@ class Sampler(nn.Module): # This conversion is necessary because FlashInfer sampling operations # return int32 (while PyTorch argmax and topk return int64). sampled = sampled.long() + _xpu_sampler_mark(sampler_timeline, "sampler_long_end_event") # Handle logprob_token_ids if specified (more efficient than full vocab) # This is used by generative_scoring API to get logprobs for specific tokens @@ -110,6 +173,7 @@ class Sampler(nn.Module): logprob_token_ids_tensors = self.gather_specific_token_logprobs( logits, sampling_metadata.logprob_token_ids, sampled ) + _xpu_sampler_mark(sampler_timeline, "sampler_specific_logprobs_end_event") if num_logprobs is None: logprobs_tensors = logprob_token_ids_tensors @@ -131,6 +195,7 @@ class Sampler(nn.Module): # Use int32 to reduce the tensor size. sampled = sampled.to(torch.int32) + _xpu_sampler_mark(sampler_timeline, "sampler_int32_end_event") # These are GPU tensors. sampler_output = SamplerOutput( @@ -140,6 +205,9 @@ class Sampler(nn.Module): sampled_token_ids=sampled.unsqueeze(-1), logprobs_tensors=logprobs_tensors, ) + _xpu_sampler_mark(sampler_timeline, "sampler_output_ready_event") + if sampler_timeline is not None: + sampler_output._xpu_sampler_timeline = sampler_timeline return sampler_output def gather_specific_token_logprobs( @@ -227,6 +295,11 @@ class Sampler(nn.Module): @staticmethod def greedy_sample(logits: torch.Tensor) -> torch.Tensor: + if ( + logits.device.type == "xpu" + and os.environ.get("VLLM_XPU_GREEDY_SAMPLE_TOPK_FALLBACK", "0") == "1" + ): + return torch.topk(logits, k=1, dim=-1).indices.view(-1) return logits.argmax(dim=-1).view(-1) def sample( @@ -246,7 +319,15 @@ class Sampler(nn.Module): if sampling_metadata.all_random: greedy_sampled = None else: + _xpu_sampler_mark( + getattr(self, "_xpu_sampler_timeline", None), + "sampler_greedy_start_event", + ) greedy_sampled = self.greedy_sample(logits) + _xpu_sampler_mark( + getattr(self, "_xpu_sampler_timeline", None), + "sampler_greedy_end_event", + ) if sampling_metadata.all_greedy: processed_logprobs = None if sampling_metadata.max_num_logprobs is not None: diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index bcab2ca2d..42fdfb843 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4,6 +4,9 @@ import functools import gc import itertools +import json +import os +import re import threading import time from collections import defaultdict @@ -56,6 +59,11 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( RoutedExpertsCapturer, ) +from vllm.model_executor.layers.fused_moe.router.route_capture import ( + pop_route_overlay_snapshot, + reset_route_overlay_snapshot, + route_overlay_enabled, +) from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( initialize_mamba_ssu_backend, ) @@ -115,6 +123,12 @@ from vllm.utils.torch_utils import ( is_quantized_kv_cache, kv_cache_dtype_str_to_dtype, ) +from vllm.utils.xpu_decode_timing import ( + begin_timing_step, + end_timing_step, + timed_region, +) +from vllm.utils.xpu_finite_trace import trace_tensor from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -218,12 +232,293 @@ if TYPE_CHECKING: logger = init_logger(__name__) + +def _select_sample_hidden_states( + hidden_states: torch.Tensor, logits_indices: torch.Tensor +) -> torch.Tensor: + if ( + os.environ.get("VLLM_XPU_SAFE_SAMPLE_HIDDEN_SELECT", "0") == "1" + and hidden_states.device.type == "xpu" + ): + if hidden_states.shape[0] == logits_indices.numel(): + return hidden_states + return torch.index_select(hidden_states, 0, logits_indices) + return hidden_states[logits_indices] + + AttnMetadataDict: TypeAlias = dict[str, AttentionMetadata] # list when ubatching is enabled PerLayerAttnMetadata: TypeAlias = list[AttnMetadataDict] | AttnMetadataDict # Wrapper for ModelRunnerOutput to support overlapped execution. +class _CompletedCopyEvent: + def synchronize(self) -> None: + return None + + +class _DeferredXpuOutputCopyEvent: + def __init__(self, output: "AsyncGPUModelRunnerOutput") -> None: + self.output = output + + def synchronize(self) -> None: + self.output.finish_deferred_xpu_output_copy() + + +def _new_copy_stream(device_type: str): + if device_type == "xpu" and hasattr(torch, "xpu"): + return torch.xpu.Stream() + return torch.cuda.Stream() + + +def _new_copy_event(device_type: str): + if device_type == "xpu" and hasattr(torch, "xpu"): + if _XPU_ASYNC_OUTPUT_DEVICE_TIMING: + return torch.xpu.Event(enable_timing=True) + return torch.xpu.Event() + if _XPU_ASYNC_OUTPUT_DEVICE_TIMING: + return torch.Event(enable_timing=True) + return torch.Event() + + +def _current_copy_stream(device_type: str): + if device_type == "xpu" and hasattr(torch, "xpu"): + return torch.xpu.current_stream() + return torch.cuda.current_stream() + + +def _copy_stream_context(stream, device_type: str): + if device_type == "xpu" and hasattr(torch, "xpu"): + return torch.xpu.stream(stream) + return torch.cuda.stream(stream) + + +_XPU_ASYNC_OUTPUT_TIMING = ( + os.environ.get("VLLM_XPU_ASYNC_OUTPUT_TIMING", "0") == "1" +) +_XPU_ASYNC_OUTPUT_TIMING_PRINT_EVERY = max( + 1, int(os.environ.get("VLLM_XPU_ASYNC_OUTPUT_TIMING_PRINT_EVERY", "1") or "1") +) +_XPU_ASYNC_OUTPUT_TIMING_SKIP_FIRST = max( + 0, int(os.environ.get("VLLM_XPU_ASYNC_OUTPUT_TIMING_SKIP_FIRST", "0") or "0") +) +_XPU_ASYNC_OUTPUT_TIMING_RANK = os.environ.get("VLLM_XPU_ASYNC_OUTPUT_TIMING_RANK") +_XPU_ASYNC_OUTPUT_SEEN = 0 +_XPU_ASYNC_OUTPUT_DEVICE_TIMING = ( + os.environ.get("VLLM_XPU_ASYNC_OUTPUT_DEVICE_TIMING", "0") == "1" +) +_XPU_PRE_SAMPLER_BOUNDARY_TIMING = ( + os.environ.get("VLLM_XPU_PRE_SAMPLER_BOUNDARY_TIMING", "0") == "1" +) +_XPU_PRE_SAMPLER_BOUNDARY_DEVICE_ELAPSED = ( + os.environ.get("VLLM_XPU_PRE_SAMPLER_BOUNDARY_DEVICE_ELAPSED", "0") == "1" +) +_XPU_PRE_SAMPLER_BOUNDARY_EVENTS = { + item.strip() + for item in os.environ.get( + "VLLM_XPU_PRE_SAMPLER_BOUNDARY_EVENTS", + ( + "forward_end_event,compute_logits_end_event," + "sample_start_event,sample_end_event,pre_async_wrap_event" + ), + ).split(",") + if item.strip() +} +_XPU_FORWARD_BOUNDARY_SYNC = ( + os.environ.get("VLLM_XPU_FORWARD_BOUNDARY_SYNC", "0") == "1" +) +_XPU_FORWARD_BOUNDARY_PRINT_EVERY = max( + 1, int(os.environ.get("VLLM_XPU_FORWARD_BOUNDARY_PRINT_EVERY", "1") or "1") +) +_XPU_FORWARD_BOUNDARY_SKIP_FIRST = max( + 0, int(os.environ.get("VLLM_XPU_FORWARD_BOUNDARY_SKIP_FIRST", "0") or "0") +) +_XPU_FORWARD_BOUNDARY_RANK = os.environ.get("VLLM_XPU_FORWARD_BOUNDARY_RANK") +_XPU_FORWARD_BOUNDARY_SEEN = 0 + + +def _xpu_async_output_rank_label() -> str: + for name in ("LOCAL_RANK", "RANK", "VLLM_LOCAL_RANK", "VLLM_RANK"): + value = os.environ.get(name) + if value is not None: + return value + return f"pid{os.getpid()}" + + +def _xpu_async_output_rank_enabled(rank: str) -> bool: + return ( + _XPU_ASYNC_OUTPUT_TIMING_RANK is None + or _XPU_ASYNC_OUTPUT_TIMING_RANK == "" + or _XPU_ASYNC_OUTPUT_TIMING_RANK == rank + ) + + +def _xpu_async_output_should_print() -> tuple[bool, int, str]: + global _XPU_ASYNC_OUTPUT_SEEN + rank = _xpu_async_output_rank_label() + if not _XPU_ASYNC_OUTPUT_TIMING or not _xpu_async_output_rank_enabled(rank): + return False, 0, rank + _XPU_ASYNC_OUTPUT_SEEN += 1 + event_id = _XPU_ASYNC_OUTPUT_SEEN + should_print = ( + event_id > _XPU_ASYNC_OUTPUT_TIMING_SKIP_FIRST + and ( + event_id - _XPU_ASYNC_OUTPUT_TIMING_SKIP_FIRST + ) + % _XPU_ASYNC_OUTPUT_TIMING_PRINT_EVERY + == 0 + ) + return should_print, event_id, rank + + +def _xpu_forward_boundary_rank_enabled(rank: str) -> bool: + return ( + _XPU_FORWARD_BOUNDARY_RANK is None + or _XPU_FORWARD_BOUNDARY_RANK == "" + or _XPU_FORWARD_BOUNDARY_RANK == rank + ) + + +def _xpu_forward_boundary_should_print() -> tuple[bool, int, str]: + global _XPU_FORWARD_BOUNDARY_SEEN + rank = _xpu_async_output_rank_label() + if ( + not _XPU_FORWARD_BOUNDARY_SYNC + or not _xpu_forward_boundary_rank_enabled(rank) + ): + return False, 0, rank + _XPU_FORWARD_BOUNDARY_SEEN += 1 + event_id = _XPU_FORWARD_BOUNDARY_SEEN + should_print = ( + event_id > _XPU_FORWARD_BOUNDARY_SKIP_FIRST + and (event_id - _XPU_FORWARD_BOUNDARY_SKIP_FIRST) + % _XPU_FORWARD_BOUNDARY_PRINT_EVERY + == 0 + ) + return should_print, event_id, rank + + +def _xpu_forward_boundary_print(payload: dict[str, object]) -> None: + print( + "[vllm-xpu-forward-boundary] " + + json.dumps(payload, sort_keys=True, separators=(",", ":")), + flush=True, + ) + + +def _xpu_async_output_ms_since(start_ns: int) -> float: + return (time.perf_counter_ns() - start_ns) / 1e6 + + +def _xpu_async_output_new_timing_event(device_type: str): + if not _XPU_ASYNC_OUTPUT_DEVICE_TIMING: + return None + try: + if device_type == "xpu" and hasattr(torch, "xpu"): + return torch.xpu.Event(enable_timing=True) + if device_type == "cuda" and torch.cuda.is_available(): + return torch.cuda.Event(enable_timing=True) + except Exception: + return None + return None + + +def _xpu_async_output_record_timing_event(device_type: str): + event = _xpu_async_output_new_timing_event(device_type) + if event is None: + return None + try: + event.record() + return event + except Exception: + return None + + +def _xpu_async_output_record_boundary_event(event_key: str, device_type: str): + is_forward_boundary = event_key in { + "forward_start_event", + "forward_end_event", + } + if not ( + _XPU_PRE_SAMPLER_BOUNDARY_TIMING + or (_XPU_FORWARD_BOUNDARY_SYNC and is_forward_boundary) + ): + return None + if ( + _XPU_PRE_SAMPLER_BOUNDARY_TIMING + and event_key not in _XPU_PRE_SAMPLER_BOUNDARY_EVENTS + and not (_XPU_FORWARD_BOUNDARY_SYNC and is_forward_boundary) + ): + return None + if ( + not _XPU_PRE_SAMPLER_BOUNDARY_TIMING + and _XPU_FORWARD_BOUNDARY_SYNC + and not is_forward_boundary + ): + return None + try: + enable_timing = _XPU_PRE_SAMPLER_BOUNDARY_DEVICE_ELAPSED + if device_type == "xpu" and hasattr(torch, "xpu"): + event = torch.xpu.Event(enable_timing=enable_timing) + elif device_type == "cuda" and torch.cuda.is_available(): + event = torch.cuda.Event(enable_timing=enable_timing) + else: + return None + event.record() + return event + except Exception: + return None + + +def _xpu_async_output_record_timeline_event(event_key: str, device_type: str): + if _XPU_ASYNC_OUTPUT_DEVICE_TIMING: + return _xpu_async_output_record_timing_event(device_type) + return _xpu_async_output_record_boundary_event(event_key, device_type) + + +def _xpu_async_output_record_timing_event_on_stream(event, stream, device_type: str): + if event is None: + return None + try: + with _copy_stream_context(stream, device_type): + event.record() + return event + except Exception: + return None + + +def _xpu_async_output_elapsed_ms(start_event, end_event) -> float | None: + if start_event is None or end_event is None: + return None + try: + return float(start_event.elapsed_time(end_event)) + except Exception: + return None + + +def _xpu_async_output_shape(tensor: torch.Tensor | None) -> list[int] | None: + if tensor is None: + return None + return [int(dim) for dim in tensor.shape] + + +def _xpu_async_output_is_pinned(tensor: torch.Tensor | None) -> bool | None: + if tensor is None or tensor.device.type != "cpu": + return None + try: + return bool(tensor.is_pinned()) + except RuntimeError: + return None + + +def _xpu_async_output_print(payload: dict[str, object]) -> None: + print( + "[vllm-xpu-async-output] " + + json.dumps(payload, sort_keys=True, separators=(",", ":")), + flush=True, + ) + + class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): def __init__( self, @@ -233,22 +528,170 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): invalid_req_indices: list[int], async_output_copy_stream: torch.cuda.Stream, vocab_size: int, + sampled_token_ids_cpu_buffer: torch.Tensor | None = None, + device_timeline: dict[str, object] | None = None, ): self._model_runner_output = model_runner_output self._invalid_req_indices = invalid_req_indices + device_type = sampled_token_ids.device.type + self._xpu_async_output_device_timeline = device_timeline or {} # Event on the copy stream so we can synchronize the non-blocking copy. - self.async_copy_ready_event = torch.Event() + self.async_copy_ready_event = _new_copy_event(device_type) # Keep a reference to the device tensor to avoid it being # deallocated until we finish copying it to the host. self._sampled_token_ids = sampled_token_ids self.vocab_size = vocab_size self._logprobs_tensors = logprobs_tensors + self._deferred_xpu_output_copy = False + self._xpu_async_output_object_id = id(self) + self._xpu_async_output_created_ns = time.perf_counter_ns() + self._xpu_async_output_copy_submit_start_ns: int | None = None + self._xpu_async_output_copy_submit_end_ns: int | None = None + self._xpu_async_output_default_before_copy_event = None + self._xpu_async_output_copy_stream_entry_event = None + self._xpu_async_output_copy_mode = "default_nonblocking" + self._xpu_async_output_copy_submit_ms: float | None = None + self._xpu_async_output_source_dtype = str(self._sampled_token_ids.dtype) + self._xpu_async_output_source_shape = _xpu_async_output_shape( + self._sampled_token_ids + ) + self._xpu_async_output_reuse_buffer_available = ( + sampled_token_ids_cpu_buffer is not None + ) + self._xpu_async_output_reuse_buffer_dtype = ( + str(sampled_token_ids_cpu_buffer.dtype) + if sampled_token_ids_cpu_buffer is not None + else None + ) + self._xpu_async_output_reuse_buffer_shape = ( + _xpu_async_output_shape(sampled_token_ids_cpu_buffer) + if sampled_token_ids_cpu_buffer is not None + else None + ) + self._xpu_async_output_reuse_buffer_reject_reason: str | None = None + + if ( + os.environ.get("VLLM_XPU_DEFER_ASYNC_OUTPUT_COPY", "0") == "1" + and self._sampled_token_ids.device.type == "xpu" + and self._logprobs_tensors is None + ): + copy_start_ns = time.perf_counter_ns() + self._xpu_async_output_copy_submit_start_ns = copy_start_ns + self._xpu_async_output_copy_mode = "deferred_xpu_copy" + self.sampled_token_ids_cpu = torch.empty( + self._sampled_token_ids.shape, + dtype=self._sampled_token_ids.dtype, + device="cpu", + ) + self._logprobs_tensors_cpu = None + self._deferred_xpu_output_copy = True + self.async_copy_ready_event = _DeferredXpuOutputCopyEvent(self) + copy_end_ns = time.perf_counter_ns() + self._xpu_async_output_copy_submit_end_ns = copy_end_ns + self._xpu_async_output_copy_submit_ms = ( + copy_end_ns - copy_start_ns + ) / 1e6 + return + + if ( + os.environ.get("VLLM_XPU_SYNC_ASYNC_OUTPUT_COPY", "0") == "1" + and self._sampled_token_ids.device.type == "xpu" + ): + copy_start_ns = time.perf_counter_ns() + self._xpu_async_output_copy_submit_start_ns = copy_start_ns + self._xpu_async_output_copy_mode = "sync_copy" + self.sampled_token_ids_cpu = self._sampled_token_ids.to("cpu") + self._logprobs_tensors_cpu = ( + LogprobsTensors( + self._logprobs_tensors.logprob_token_ids.to("cpu"), + self._logprobs_tensors.logprobs.to("cpu"), + self._logprobs_tensors.selected_token_ranks.to("cpu"), + self._logprobs_tensors.cu_num_generated_tokens, + ) + if self._logprobs_tensors + else None + ) + self.async_copy_ready_event = _CompletedCopyEvent() + copy_end_ns = time.perf_counter_ns() + self._xpu_async_output_copy_submit_end_ns = copy_end_ns + self._xpu_async_output_copy_submit_ms = ( + copy_end_ns - copy_start_ns + ) / 1e6 + return + + reuse_requested = ( + os.environ.get("VLLM_XPU_REUSE_ASYNC_OUTPUT_COPY_BUFFER", "0") == "1" + ) + reuse_allowed = True + if not reuse_requested: + reuse_allowed = False + self._xpu_async_output_reuse_buffer_reject_reason = "disabled" + elif self._sampled_token_ids.device.type != "xpu": + reuse_allowed = False + self._xpu_async_output_reuse_buffer_reject_reason = "not_xpu" + elif self._logprobs_tensors is not None: + reuse_allowed = False + self._xpu_async_output_reuse_buffer_reject_reason = "has_logprobs" + elif sampled_token_ids_cpu_buffer is None: + reuse_allowed = False + self._xpu_async_output_reuse_buffer_reject_reason = "no_buffer" + elif sampled_token_ids_cpu_buffer.dtype != self._sampled_token_ids.dtype: + reuse_allowed = False + self._xpu_async_output_reuse_buffer_reject_reason = "dtype_mismatch" + elif ( + sampled_token_ids_cpu_buffer.shape[0] < self._sampled_token_ids.shape[0] + or sampled_token_ids_cpu_buffer.shape[1] < self._sampled_token_ids.shape[1] + ): + reuse_allowed = False + self._xpu_async_output_reuse_buffer_reject_reason = "shape_mismatch" + + if reuse_allowed and sampled_token_ids_cpu_buffer is not None: + copy_start_ns = time.perf_counter_ns() + self._xpu_async_output_copy_submit_start_ns = copy_start_ns + self._xpu_async_output_copy_mode = "reuse_buffer_nonblocking" + self.sampled_token_ids_cpu = sampled_token_ids_cpu_buffer[ + : self._sampled_token_ids.shape[0], + : self._sampled_token_ids.shape[1], + ] + self._logprobs_tensors_cpu = None + default_stream = _current_copy_stream(device_type) + self._xpu_async_output_default_before_copy_event = ( + _xpu_async_output_record_timing_event(device_type) + ) + self._xpu_async_output_copy_stream_entry_event = ( + _xpu_async_output_new_timing_event(device_type) + ) + with _copy_stream_context(async_output_copy_stream, device_type): + if self._xpu_async_output_copy_stream_entry_event is not None: + self._xpu_async_output_copy_stream_entry_event.record() + async_output_copy_stream.wait_stream(default_stream) + self.sampled_token_ids_cpu.copy_( + self._sampled_token_ids, non_blocking=True + ) + self.async_copy_ready_event.record() + copy_end_ns = time.perf_counter_ns() + self._xpu_async_output_copy_submit_end_ns = copy_end_ns + self._xpu_async_output_copy_submit_ms = ( + copy_end_ns - copy_start_ns + ) / 1e6 + return # Initiate the copy on a separate stream, but do not synchronize it. - default_stream = torch.cuda.current_stream() - with torch.cuda.stream(async_output_copy_stream): + copy_start_ns = time.perf_counter_ns() + self._xpu_async_output_copy_submit_start_ns = copy_start_ns + self._xpu_async_output_copy_mode = "default_nonblocking" + default_stream = _current_copy_stream(device_type) + self._xpu_async_output_default_before_copy_event = ( + _xpu_async_output_record_timing_event(device_type) + ) + self._xpu_async_output_copy_stream_entry_event = ( + _xpu_async_output_new_timing_event(device_type) + ) + with _copy_stream_context(async_output_copy_stream, device_type): + if self._xpu_async_output_copy_stream_entry_event is not None: + self._xpu_async_output_copy_stream_entry_event.record() async_output_copy_stream.wait_stream(default_stream) self.sampled_token_ids_cpu = self._sampled_token_ids.to( "cpu", non_blocking=True @@ -259,36 +702,517 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): else None ) self.async_copy_ready_event.record() + copy_end_ns = time.perf_counter_ns() + self._xpu_async_output_copy_submit_end_ns = copy_end_ns + self._xpu_async_output_copy_submit_ms = (copy_end_ns - copy_start_ns) / 1e6 + + def finish_deferred_xpu_output_copy(self) -> None: + if not self._deferred_xpu_output_copy: + return + with torch.inference_mode(): + if os.environ.get("VLLM_XPU_DEFERRED_COPY_SYNC_DEVICE", "1") != "0": + torch.xpu.synchronize() + self.sampled_token_ids_cpu.copy_(self._sampled_token_ids) + self._deferred_xpu_output_copy = False def get_output(self) -> ModelRunnerOutput: """Copy the device tensors to the host and return a ModelRunnerOutput. This function blocks until the copy is finished. """ + should_print, event_id, rank = _xpu_async_output_should_print() + total_start_ns = time.perf_counter_ns() + created_ns = getattr(self, "_xpu_async_output_created_ns", None) + copy_submit_start_ns = getattr( + self, "_xpu_async_output_copy_submit_start_ns", None + ) + copy_submit_end_ns = getattr( + self, "_xpu_async_output_copy_submit_end_ns", None + ) + created_to_get_output_ms = ( + (total_start_ns - created_ns) / 1e6 + if isinstance(created_ns, int) + else None + ) + copy_start_to_get_output_ms = ( + (total_start_ns - copy_submit_start_ns) / 1e6 + if isinstance(copy_submit_start_ns, int) + else None + ) + copy_end_to_get_output_ms = ( + (total_start_ns - copy_submit_end_ns) / 1e6 + if isinstance(copy_submit_end_ns, int) + else None + ) max_gen_len = self.sampled_token_ids_cpu.shape[-1] - self.async_copy_ready_event.synchronize() + sync_start_ns = time.perf_counter_ns() + device_timeline = getattr(self, "_xpu_async_output_device_timeline", {}) + sync_stage_split_enabled = ( + os.environ.get("VLLM_XPU_ASYNC_OUTPUT_SYNC_STAGE_SPLIT", "0") == "1" + and isinstance(device_timeline, dict) + ) + pre_sampler_stage_sync_ms: dict[str, float] | None = None + pre_sampler_device_elapsed_ms: dict[str, float | None] | None = None + pre_sampler_metadata: dict[str, object] | None = None + sync_pre_sampler_split_enabled = ( + os.environ.get("VLLM_XPU_ASYNC_OUTPUT_SYNC_PRE_SAMPLER_SPLIT", "0") + == "1" + and isinstance(device_timeline, dict) + ) + stage_sample_end_sync_ms = None + stage_state_update_sync_ms = None + stage_bookkeeping_sync_ms = None + stage_pre_async_wrap_sync_ms = None + sampler_stage_sync_ms: dict[str, float] | None = None + sampler_device_elapsed_ms: dict[str, float | None] | None = None + sampler_metadata: dict[str, object] | None = None + sync_sampler_split_enabled = ( + os.environ.get("VLLM_XPU_ASYNC_OUTPUT_SYNC_SAMPLER_SPLIT", "0") == "1" + and isinstance(device_timeline, dict) + ) + sync_split_enabled = ( + ( + os.environ.get("VLLM_XPU_ASYNC_OUTPUT_SYNC_SPLIT", "0") == "1" + or sync_stage_split_enabled + or sync_sampler_split_enabled + or sync_pre_sampler_split_enabled + ) + and ( + self._xpu_async_output_default_before_copy_event is not None + or sync_pre_sampler_split_enabled + or sync_sampler_split_enabled + or sync_stage_split_enabled + ) + ) + default_ready_sync_ms = None + copy_after_default_sync_ms = None + if sync_split_enabled: + if sync_pre_sampler_split_enabled: + pre_sampler_events = ( + ("execute_entry_event", "execute_entry"), + ("forward_start_event", "forward_start"), + ("forward_end_event", "forward_end"), + ("select_sample_hidden_start_event", "select_sample_hidden_start"), + ("select_sample_hidden_end_event", "select_sample_hidden_end"), + ("clone_sample_hidden_start_event", "clone_sample_hidden_start"), + ("clone_sample_hidden_end_event", "clone_sample_hidden_end"), + ("compute_logits_start_event", "compute_logits_start"), + ("compute_logits_end_event", "compute_logits_end"), + ("local_argmax_start_event", "local_argmax_start"), + ("local_argmax_end_event", "local_argmax_end"), + ("execute_state_ready_event", "execute_state_ready"), + ("sample_start_event", "sample_start"), + ("sampler_entry_event", "sampler_entry"), + ) + pre_sampler_stage_sync_ms = {} + for event_key, timing_key in pre_sampler_events: + event = device_timeline.get(event_key) + if event is None: + continue + pre_sampler_sync_start_ns = time.perf_counter_ns() + event.synchronize() + pre_sampler_stage_sync_ms[timing_key] = ( + _xpu_async_output_ms_since(pre_sampler_sync_start_ns) + ) + if sync_sampler_split_enabled: + sampler_events = ( + ("sampler_entry_event", "sampler_entry"), + ("sampler_raw_logprobs_end_event", "sampler_raw_logprobs"), + ("sampler_fp32_cast_end_event", "sampler_fp32_cast"), + ("sampler_processors_end_event", "sampler_processors"), + ("sampler_greedy_start_event", "sampler_greedy_start"), + ("sampler_greedy_end_event", "sampler_greedy_end"), + ("sampler_sample_call_end_event", "sampler_sample_call"), + ("sampler_long_end_event", "sampler_long"), + ( + "sampler_specific_logprobs_end_event", + "sampler_specific_logprobs", + ), + ("sampler_int32_end_event", "sampler_int32"), + ("sampler_output_ready_event", "sampler_output_ready"), + ) + sampler_stage_sync_ms = {} + for event_key, timing_key in sampler_events: + event = device_timeline.get(event_key) + if event is None: + continue + sampler_sync_start_ns = time.perf_counter_ns() + event.synchronize() + sampler_stage_sync_ms[timing_key] = _xpu_async_output_ms_since( + sampler_sync_start_ns + ) + if sync_stage_split_enabled: + stage_events = ( + ("sample_end_event", "sample"), + ("state_update_end_event", "state_update"), + ("bookkeeping_end_event", "bookkeeping"), + ("pre_async_wrap_event", "pre_async_wrap"), + ) + stage_timings: dict[str, float] = {} + for event_key, timing_key in stage_events: + event = device_timeline.get(event_key) + if event is None: + continue + stage_sync_start_ns = time.perf_counter_ns() + event.synchronize() + stage_timings[timing_key] = _xpu_async_output_ms_since( + stage_sync_start_ns + ) + stage_sample_end_sync_ms = stage_timings.get("sample") + stage_state_update_sync_ms = stage_timings.get("state_update") + stage_bookkeeping_sync_ms = stage_timings.get("bookkeeping") + stage_pre_async_wrap_sync_ms = stage_timings.get("pre_async_wrap") + if self._xpu_async_output_default_before_copy_event is not None: + default_sync_start_ns = time.perf_counter_ns() + self._xpu_async_output_default_before_copy_event.synchronize() + default_ready_sync_ms = _xpu_async_output_ms_since( + default_sync_start_ns + ) + copy_sync_start_ns = time.perf_counter_ns() + self.async_copy_ready_event.synchronize() + copy_after_default_sync_ms = _xpu_async_output_ms_since( + copy_sync_start_ns + ) + else: + self.async_copy_ready_event.synchronize() + sync_end_ns = time.perf_counter_ns() + sync_ms = _xpu_async_output_ms_since(sync_start_ns) + copy_end_to_sync_done_ms = ( + (sync_end_ns - copy_submit_end_ns) / 1e6 + if isinstance(copy_submit_end_ns, int) + else None + ) + device_sample_ms = _xpu_async_output_elapsed_ms( + device_timeline.get("sample_start_event") + if isinstance(device_timeline, dict) + else None, + device_timeline.get("sample_end_event") + if isinstance(device_timeline, dict) + else None, + ) + device_sample_end_to_state_update_end_ms = _xpu_async_output_elapsed_ms( + device_timeline.get("sample_end_event") + if isinstance(device_timeline, dict) + else None, + device_timeline.get("state_update_end_event") + if isinstance(device_timeline, dict) + else None, + ) + device_state_update_to_bookkeeping_end_ms = _xpu_async_output_elapsed_ms( + device_timeline.get("state_update_end_event") + if isinstance(device_timeline, dict) + else None, + device_timeline.get("bookkeeping_end_event") + if isinstance(device_timeline, dict) + else None, + ) + device_bookkeeping_to_pre_async_wrap_ms = _xpu_async_output_elapsed_ms( + device_timeline.get("bookkeeping_end_event") + if isinstance(device_timeline, dict) + else None, + device_timeline.get("pre_async_wrap_event") + if isinstance(device_timeline, dict) + else None, + ) + device_pre_async_wrap_to_default_before_copy_ms = ( + _xpu_async_output_elapsed_ms( + device_timeline.get("pre_async_wrap_event") + if isinstance(device_timeline, dict) + else None, + self._xpu_async_output_default_before_copy_event, + ) + ) + device_copy_stream_entry_to_ready_ms = _xpu_async_output_elapsed_ms( + self._xpu_async_output_copy_stream_entry_event, + self.async_copy_ready_event, + ) + device_default_before_copy_to_ready_ms = _xpu_async_output_elapsed_ms( + self._xpu_async_output_default_before_copy_event, + self.async_copy_ready_event, + ) + device_sample_start_to_copy_ready_ms = _xpu_async_output_elapsed_ms( + device_timeline.get("sample_start_event") + if isinstance(device_timeline, dict) + else None, + self.async_copy_ready_event, + ) + device_sample_end_to_copy_ready_ms = _xpu_async_output_elapsed_ms( + device_timeline.get("sample_end_event") + if isinstance(device_timeline, dict) + else None, + self.async_copy_ready_event, + ) + device_pre_async_wrap_to_copy_ready_ms = _xpu_async_output_elapsed_ms( + device_timeline.get("pre_async_wrap_event") + if isinstance(device_timeline, dict) + else None, + self.async_copy_ready_event, + ) + if isinstance(device_timeline, dict): + pre_sampler_metadata = { + key: value + for key, value in device_timeline.items() + if key.startswith("pre_sampler_") and not key.endswith("_event") + } + pre_sampler_device_elapsed_ms = { + "execute_entry_to_forward_start": _xpu_async_output_elapsed_ms( + device_timeline.get("execute_entry_event"), + device_timeline.get("forward_start_event"), + ), + "forward_start_to_forward_end": _xpu_async_output_elapsed_ms( + device_timeline.get("forward_start_event"), + device_timeline.get("forward_end_event"), + ), + "forward_end_to_select_start": _xpu_async_output_elapsed_ms( + device_timeline.get("forward_end_event"), + device_timeline.get("select_sample_hidden_start_event"), + ), + "select_start_to_select_end": _xpu_async_output_elapsed_ms( + device_timeline.get("select_sample_hidden_start_event"), + device_timeline.get("select_sample_hidden_end_event"), + ), + "select_end_to_logits_start": _xpu_async_output_elapsed_ms( + device_timeline.get("select_sample_hidden_end_event"), + device_timeline.get("compute_logits_start_event"), + ), + "logits_start_to_logits_end": _xpu_async_output_elapsed_ms( + device_timeline.get("compute_logits_start_event"), + device_timeline.get("compute_logits_end_event"), + ), + "logits_end_to_execute_state_ready": _xpu_async_output_elapsed_ms( + device_timeline.get("compute_logits_end_event"), + device_timeline.get("execute_state_ready_event"), + ), + "execute_entry_to_logits_end": _xpu_async_output_elapsed_ms( + device_timeline.get("execute_entry_event"), + device_timeline.get("compute_logits_end_event"), + ), + "logits_end_to_sample_start": _xpu_async_output_elapsed_ms( + device_timeline.get("compute_logits_end_event"), + device_timeline.get("sample_start_event"), + ), + "logits_end_to_sampler_entry": _xpu_async_output_elapsed_ms( + device_timeline.get("compute_logits_end_event"), + device_timeline.get("sampler_entry_event"), + ), + "execute_state_ready_to_sample_start": ( + _xpu_async_output_elapsed_ms( + device_timeline.get("execute_state_ready_event"), + device_timeline.get("sample_start_event"), + ) + ), + "sample_start_to_sampler_entry": _xpu_async_output_elapsed_ms( + device_timeline.get("sample_start_event"), + device_timeline.get("sampler_entry_event"), + ), + "forward_start_to_sampler_entry": _xpu_async_output_elapsed_ms( + device_timeline.get("forward_start_event"), + device_timeline.get("sampler_entry_event"), + ), + } + sampler_metadata = { + key: value + for key, value in device_timeline.items() + if key.startswith("sampler_") and not key.endswith("_event") + } + sampler_device_elapsed_ms = { + "entry_to_raw_logprobs": _xpu_async_output_elapsed_ms( + device_timeline.get("sampler_entry_event"), + device_timeline.get("sampler_raw_logprobs_end_event"), + ), + "raw_logprobs_to_fp32_cast": _xpu_async_output_elapsed_ms( + device_timeline.get("sampler_raw_logprobs_end_event"), + device_timeline.get("sampler_fp32_cast_end_event"), + ), + "fp32_cast_to_processors": _xpu_async_output_elapsed_ms( + device_timeline.get("sampler_fp32_cast_end_event"), + device_timeline.get("sampler_processors_end_event"), + ), + "processors_to_greedy_start": _xpu_async_output_elapsed_ms( + device_timeline.get("sampler_processors_end_event"), + device_timeline.get("sampler_greedy_start_event"), + ), + "greedy_start_to_greedy_end": _xpu_async_output_elapsed_ms( + device_timeline.get("sampler_greedy_start_event"), + device_timeline.get("sampler_greedy_end_event"), + ), + "greedy_end_to_sample_call": _xpu_async_output_elapsed_ms( + device_timeline.get("sampler_greedy_end_event"), + device_timeline.get("sampler_sample_call_end_event"), + ), + "sample_call_to_long": _xpu_async_output_elapsed_ms( + device_timeline.get("sampler_sample_call_end_event"), + device_timeline.get("sampler_long_end_event"), + ), + "long_to_specific_logprobs": _xpu_async_output_elapsed_ms( + device_timeline.get("sampler_long_end_event"), + device_timeline.get("sampler_specific_logprobs_end_event"), + ), + "specific_logprobs_to_int32": _xpu_async_output_elapsed_ms( + device_timeline.get("sampler_specific_logprobs_end_event"), + device_timeline.get("sampler_int32_end_event"), + ), + "int32_to_output_ready": _xpu_async_output_elapsed_ms( + device_timeline.get("sampler_int32_end_event"), + device_timeline.get("sampler_output_ready_event"), + ), + "entry_to_output_ready": _xpu_async_output_elapsed_ms( + device_timeline.get("sampler_entry_event"), + device_timeline.get("sampler_output_ready_event"), + ), + } # Release the device tensors once the copy has completed. + release_start_ns = time.perf_counter_ns() del self._logprobs_tensors del self._sampled_token_ids + release_ms = _xpu_async_output_ms_since(release_start_ns) + token_list_ms: float | None = None + logprobs_ms: float | None = None + parse_output_ms: float | None = None + invalid_clear_ms: float | None = None + list_mode = "none" if max_gen_len == 1: - valid_sampled_token_ids = self.sampled_token_ids_cpu.tolist() + if ( + os.environ.get("VLLM_XPU_FAST_ASYNC_OUTPUT_LIST", "0") == "1" + and self.sampled_token_ids_cpu.shape[0] == 1 + and self._logprobs_tensors_cpu is None + and not self._invalid_req_indices + ): + list_mode = "fast_scalar" + list_start_ns = time.perf_counter_ns() + with timed_region("gpu_model_runner.async_output_fast_list"): + valid_sampled_token_ids = [ + [int(self.sampled_token_ids_cpu[0, 0])] + ] + token_list_ms = _xpu_async_output_ms_since(list_start_ns) + else: + list_mode = "tolist" + list_start_ns = time.perf_counter_ns() + with timed_region("gpu_model_runner.async_output_tolist"): + valid_sampled_token_ids = self.sampled_token_ids_cpu.tolist() + token_list_ms = _xpu_async_output_ms_since(list_start_ns) + invalid_start_ns = time.perf_counter_ns() for i in self._invalid_req_indices: valid_sampled_token_ids[i].clear() + invalid_clear_ms = _xpu_async_output_ms_since(invalid_start_ns) logprobs_lists = None if self._logprobs_tensors_cpu is not None: + logprobs_start_ns = time.perf_counter_ns() logprobs_lists = self._logprobs_tensors_cpu.tolists() + logprobs_ms = _xpu_async_output_ms_since(logprobs_start_ns) else: + list_mode = "rejection_sampler_parse" + parse_start_ns = time.perf_counter_ns() valid_sampled_token_ids, logprobs_lists = RejectionSampler.parse_output( self.sampled_token_ids_cpu, self.vocab_size, self._invalid_req_indices, logprobs_tensors=self._logprobs_tensors_cpu, ) + parse_output_ms = _xpu_async_output_ms_since(parse_start_ns) + assign_start_ns = time.perf_counter_ns() output = self._model_runner_output output.sampled_token_ids = valid_sampled_token_ids output.logprobs = logprobs_lists + assign_ms = _xpu_async_output_ms_since(assign_start_ns) + total_ms = _xpu_async_output_ms_since(total_start_ns) + if should_print: + _xpu_async_output_print( + { + "event": "get_output", + "event_id": event_id, + "rank": rank, + "object_id": getattr( + self, "_xpu_async_output_object_id", None + ), + "copy_mode": self._xpu_async_output_copy_mode, + "copy_submit_ms": self._xpu_async_output_copy_submit_ms, + "created_to_get_output_ms": created_to_get_output_ms, + "copy_start_to_get_output_ms": copy_start_to_get_output_ms, + "copy_end_to_get_output_ms": copy_end_to_get_output_ms, + "copy_end_to_sync_done_ms": copy_end_to_sync_done_ms, + "sync_pre_sampler_split_enabled": ( + sync_pre_sampler_split_enabled + ), + "sync_stage_split_enabled": sync_stage_split_enabled, + "sync_sampler_split_enabled": sync_sampler_split_enabled, + "sync_split_enabled": sync_split_enabled, + "pre_sampler_stage_sync_ms": pre_sampler_stage_sync_ms, + "pre_sampler_device_elapsed_ms": pre_sampler_device_elapsed_ms, + "pre_sampler_metadata": pre_sampler_metadata, + "sampler_stage_sync_ms": sampler_stage_sync_ms, + "sampler_device_elapsed_ms": sampler_device_elapsed_ms, + "sampler_metadata": sampler_metadata, + "stage_sample_end_sync_ms": stage_sample_end_sync_ms, + "stage_state_update_sync_ms": stage_state_update_sync_ms, + "stage_bookkeeping_sync_ms": stage_bookkeeping_sync_ms, + "stage_pre_async_wrap_sync_ms": stage_pre_async_wrap_sync_ms, + "default_ready_sync_ms": default_ready_sync_ms, + "copy_after_default_sync_ms": copy_after_default_sync_ms, + "device_sample_ms": device_sample_ms, + "device_sample_end_to_state_update_end_ms": ( + device_sample_end_to_state_update_end_ms + ), + "device_state_update_to_bookkeeping_end_ms": ( + device_state_update_to_bookkeeping_end_ms + ), + "device_bookkeeping_to_pre_async_wrap_ms": ( + device_bookkeeping_to_pre_async_wrap_ms + ), + "device_pre_async_wrap_to_default_before_copy_ms": ( + device_pre_async_wrap_to_default_before_copy_ms + ), + "device_copy_stream_entry_to_ready_ms": ( + device_copy_stream_entry_to_ready_ms + ), + "device_default_before_copy_to_ready_ms": ( + device_default_before_copy_to_ready_ms + ), + "device_sample_start_to_copy_ready_ms": ( + device_sample_start_to_copy_ready_ms + ), + "device_sample_end_to_copy_ready_ms": ( + device_sample_end_to_copy_ready_ms + ), + "device_pre_async_wrap_to_copy_ready_ms": ( + device_pre_async_wrap_to_copy_ready_ms + ), + "sync_ms": sync_ms, + "release_ms": release_ms, + "token_list_ms": token_list_ms, + "invalid_clear_ms": invalid_clear_ms, + "logprobs_ms": logprobs_ms, + "parse_output_ms": parse_output_ms, + "assign_ms": assign_ms, + "total_ms": total_ms, + "list_mode": list_mode, + "max_gen_len": int(max_gen_len), + "invalid_req_count": len(self._invalid_req_indices), + "has_logprobs": self._logprobs_tensors_cpu is not None, + "source_dtype": getattr( + self, "_xpu_async_output_source_dtype", None + ), + "source_shape": getattr( + self, "_xpu_async_output_source_shape", None + ), + "cpu_dtype": str(self.sampled_token_ids_cpu.dtype), + "cpu_shape": _xpu_async_output_shape(self.sampled_token_ids_cpu), + "cpu_is_pinned": _xpu_async_output_is_pinned( + self.sampled_token_ids_cpu + ), + "reuse_buffer_available": ( + self._xpu_async_output_reuse_buffer_available + ), + "reuse_buffer_dtype": self._xpu_async_output_reuse_buffer_dtype, + "reuse_buffer_shape": self._xpu_async_output_reuse_buffer_shape, + "reuse_buffer_reject_reason": ( + self._xpu_async_output_reuse_buffer_reject_reason + ), + "event_class": type(self.async_copy_ready_event).__name__, + } + ) return output @@ -346,17 +1270,20 @@ class AsyncGPUPoolingModelRunnerOutput(AsyncModelRunnerOutput): async_output_copy_stream: torch.cuda.Stream, ): self._model_runner_output = model_runner_output + device_type = raw_pooler_output.device.type if isinstance( + raw_pooler_output, torch.Tensor + ) else "cuda" # Event on the copy stream so we can synchronize the non-blocking copy. - self.async_copy_ready_event = torch.Event() + self.async_copy_ready_event = _new_copy_event(device_type) # Keep a reference to the device tensors to avoid them being # deallocated until we finish copying it to the host. self._raw_pooler_output = raw_pooler_output # Initiate the copy on a separate stream, but do not synchronize it. - default_stream = torch.cuda.current_stream() - with torch.cuda.stream(async_output_copy_stream): + default_stream = _current_copy_stream(device_type) + with _copy_stream_context(async_output_copy_stream, device_type): async_output_copy_stream.wait_stream(default_stream) self._model_runner_output.pooler_output = _copy_pooler_output_to_cpu( raw_pooler_output=self._raw_pooler_output, @@ -389,6 +1316,7 @@ class ExecuteModelState(NamedTuple): ec_connector_output: ECConnectorOutput | None cudagraph_stats: CUDAGraphStat | None slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None + device_timeline: dict[str, object] | None class GPUModelRunner( @@ -440,6 +1368,49 @@ class GPUModelRunner( self.calculate_kv_scales = self.cache_config.calculate_kv_scales self.dcp_world_size = self.parallel_config.decode_context_parallel_size self.dcp_rank = 0 if self.dcp_world_size <= 1 else get_dcp_group().rank_in_group + self.xpu_model_input_trace_file = os.environ.get( + "VLLM_XPU_MODEL_INPUT_TRACE_FILE" + ) + self.xpu_model_input_trace_max_lines = int( + os.environ.get("VLLM_XPU_MODEL_INPUT_TRACE_MAX_LINES", "0") or "0" + ) + self.xpu_model_input_trace_lines = 0 + self.xpu_model_input_trace_rank = os.environ.get( + "VLLM_XPU_MODEL_INPUT_TRACE_RANK", "" + ).strip() + self.xpu_replay_microscope_file = os.environ.get( + "VLLM_XPU_REPLAY_MICROSCOPE_FILE" + ) + self.xpu_replay_microscope_max_lines = int( + os.environ.get("VLLM_XPU_REPLAY_MICROSCOPE_MAX_LINES", "0") or "0" + ) + self.xpu_replay_microscope_lines = 0 + self.xpu_replay_microscope_rank = os.environ.get( + "VLLM_XPU_REPLAY_MICROSCOPE_RANK", "" + ).strip() + self.xpu_replay_microscope_req_regex = os.environ.get( + "VLLM_XPU_REPLAY_MICROSCOPE_REQ_REGEX", "" + ).strip() + self.xpu_replay_microscope_tensor_limit = int( + os.environ.get("VLLM_XPU_REPLAY_MICROSCOPE_TENSOR_LIMIT", "8") or "8" + ) + self.xpu_replay_microscope_topk = int( + os.environ.get("VLLM_XPU_REPLAY_MICROSCOPE_TOPK", "8") or "8" + ) + self.xpu_cow_worker_trace_file = os.environ.get( + "VLLM_XPU_COW_WORKER_TRACE_FILE" + ) + self.xpu_cow_worker_trace_max_lines = int( + os.environ.get("VLLM_XPU_COW_WORKER_TRACE_MAX_LINES", "0") or "0" + ) + self.xpu_cow_worker_trace_lines = 0 + self.xpu_cow_worker_trace_rank = os.environ.get( + "VLLM_XPU_COW_WORKER_TRACE_RANK", "" + ).strip() + try: + self.tp_rank = get_tp_group().rank_in_group + except Exception: + self.tp_rank = int(os.environ.get("LOCAL_RANK", "0") or "0") self.max_num_tokens = scheduler_config.max_num_batched_tokens self.max_num_reqs = scheduler_config.max_num_seqs @@ -655,8 +1626,9 @@ class GPUModelRunner( # when async scheduling is enabled. self.prepare_inputs_event: torch.Event | None = None if self.use_async_scheduling: - self.async_output_copy_stream = torch.cuda.Stream() - self.prepare_inputs_event = torch.Event() + device_type = self.device.type + self.async_output_copy_stream = _new_copy_stream(device_type) + self.prepare_inputs_event = _new_copy_event(device_type) # self.cudagraph_batch_sizes sorts in ascending order. if ( @@ -777,6 +1749,8 @@ class GPUModelRunner( # Cudagraph dispatcher for runtime cudagraph dispatching. self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config) + self.xpu_decode_cudagraph_replay_count = 0 + self.xpu_cudagraph_last_prefill_clear_req_key = "" self.mm_budget = ( MultiModalBudget(self.vllm_config, self.mm_registry) @@ -816,6 +1790,22 @@ class GPUModelRunner( device="cpu", pin_memory=self.pin_memory, ) + self.async_sampled_token_ids_pinned_cpu_buffers: list[torch.Tensor] = [] + self.async_sampled_token_ids_pinned_cpu_buffer_index = 0 + if os.environ.get("VLLM_XPU_REUSE_ASYNC_OUTPUT_COPY_BUFFER", "0") == "1": + num_async_output_buffers = max( + 3, + int(os.environ.get("VLLM_XPU_ASYNC_OUTPUT_COPY_BUFFER_SLOTS", "3")), + ) + self.async_sampled_token_ids_pinned_cpu_buffers = [ + torch.empty( + (self.max_num_reqs, 1), + dtype=torch.int32, + device="cpu", + pin_memory=self.pin_memory, + ) + for _ in range(num_async_output_buffers) + ] # Pre-allocated tensor for copying valid sampled token counts to CPU, # with dedicated stream for overlapping and event for coordination. @@ -1058,7 +2048,7 @@ class GPUModelRunner( def _get_or_create_async_output_copy_stream(self) -> torch.cuda.Stream: stream = self.async_output_copy_stream if stream is None: - stream = torch.cuda.Stream() + stream = _new_copy_stream(self.device.type) self.async_output_copy_stream = stream return stream @@ -1123,7 +2113,10 @@ class GPUModelRunner( self.speculative_config is not None and self.speculative_config.use_ngram_gpu() ) - if is_ngram_gpu: + has_ngram_gpu_buffers = is_ngram_gpu and hasattr( + self, "token_ids_gpu_tensor" + ) + if has_ngram_gpu_buffers: ngram_gpu_new_reqs: list[CachedRequestState] = [] reqs_to_add: list[CachedRequestState] = [] @@ -1193,13 +2186,14 @@ class GPUModelRunner( reqs_to_add.append(req_state) # Track new requests for ngram_gpu full tensor copy - if is_ngram_gpu: + if has_ngram_gpu_buffers: ngram_gpu_new_reqs.append(req_state) # Update the states of the running/resumed requests. is_last_rank = get_pp_group().is_last_rank req_data = scheduler_output.scheduled_cached_reqs scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens + trace_cow_worker = self._xpu_cow_worker_trace_enabled() # Save scheduler-allocated spec lengths before trimming so # prev_num_draft_len keeps the optimistic count for rejection correction. @@ -1208,6 +2202,9 @@ class GPUModelRunner( self.speculative_config is not None and self.speculative_config.use_ngram_gpu() ): + trace_ngram_pp = os.environ.get("VLLM_XPU_TRACE_NGRAM_PP") + original_total = scheduler_output.total_num_scheduled_tokens + original_num_scheduled = scheduler_output.num_scheduled_tokens.copy() for req_id, toks in scheduled_spec_tokens.items(): original_num_spec_per_req[req_id] = len(toks) update_scheduler_for_invalid_drafts( @@ -1216,6 +2213,34 @@ class GPUModelRunner( scheduler_output, self.input_batch.req_id_to_index, ) + if trace_ngram_pp: + valid_counts: dict[str, int] = {} + if self._num_valid_draft_tokens_cpu is not None: + for req_id in req_data.req_ids: + req_index = self.input_batch.req_id_to_index.get(req_id) + if req_index is not None: + valid_counts[req_id] = int( + self._num_valid_draft_tokens_cpu[req_index].item() + ) + logger.warning( + "XPU ngram trim pp_rank=%s first=%s last=%s " + "total=%s->%s scheduled=%s->%s spec_lens=%s->%s " + "valid_counts=%s cached_req_ids=%s", + get_pp_group().rank_in_group, + get_pp_group().is_first_rank, + get_pp_group().is_last_rank, + original_total, + scheduler_output.total_num_scheduled_tokens, + original_num_scheduled, + scheduler_output.num_scheduled_tokens, + original_num_spec_per_req, + { + req_id: len(toks) + for req_id, toks in scheduled_spec_tokens.items() + }, + valid_counts, + req_data.req_ids, + ) if self.use_async_spec_decode: self.prev_num_draft_tokens.np.fill(0) @@ -1226,6 +2251,30 @@ class GPUModelRunner( resumed_from_preemption = req_id in req_data.resumed_req_ids num_output_tokens = req_data.num_output_tokens[i] req_index = self.input_batch.req_id_to_index.get(req_id) + if trace_cow_worker: + self._trace_xpu_cow_worker_event( + stage="before_cached_req_update", + scheduler_output=scheduler_output, + req_id=req_id, + req_state=req_state, + req_index=req_index, + extra={ + "scheduled_cached_index": int(i), + "scheduler_num_computed_tokens": int( + num_computed_tokens + ), + "scheduler_num_output_tokens": int(num_output_tokens), + "resumed_from_preemption": bool( + resumed_from_preemption + ), + "new_block_ids_lens": None + if new_block_ids is None + else [len(ids) for ids in new_block_ids], + "scheduled_spec_len": len( + scheduled_spec_tokens.get(req_id, ()) + ), + }, + ) if req_state.prev_num_draft_len and self.use_async_scheduling: # prev_num_draft_len is used in async scheduling mode with @@ -1264,13 +2313,30 @@ class GPUModelRunner( optimistic_num_accepted ) - if is_ngram_gpu and optimistic_num_accepted > 0: + if has_ngram_gpu_buffers and optimistic_num_accepted > 0: self.input_batch.num_tokens_no_spec[req_index] += ( optimistic_num_accepted ) # Update the cached states. req_state.num_computed_tokens = num_computed_tokens + if trace_cow_worker: + self._trace_xpu_cow_worker_event( + stage="after_cached_req_counter_update", + scheduler_output=scheduler_output, + req_id=req_id, + req_state=req_state, + req_index=req_index, + extra={ + "scheduled_cached_index": int(i), + "scheduler_num_computed_tokens": int( + num_computed_tokens + ), + "scheduled_spec_len": len( + scheduled_spec_tokens.get(req_id, ()) + ), + }, + ) if not is_last_rank: if not req_data.new_token_ids: @@ -1286,13 +2352,32 @@ class GPUModelRunner( num_new_tokens = ( num_computed_tokens + len(new_token_ids) - req_state.num_tokens ) - if num_new_tokens == 1: - # Avoid slicing list in most common case. - req_state.output_token_ids.append(new_token_ids[-1]) - elif num_new_tokens > 0: - req_state.output_token_ids.extend( - new_token_ids[-num_new_tokens:] - ) + if num_new_tokens > 0: + if os.environ.get("VLLM_XPU_TRACE_NGRAM_PP") and ( + num_new_tokens > len(new_token_ids) + ): + logger.warning( + "XPU PP state update has fewer token ids than " + "computed delta pp_rank=%s req_id=%s " + "num_new_tokens=%s new_token_ids=%s " + "num_computed_tokens=%s req_state_num_tokens=%s " + "scheduled=%s spec_lens=%s", + get_pp_group().rank_in_group, + req_id, + num_new_tokens, + new_token_ids, + num_computed_tokens, + req_state.num_tokens, + scheduler_output.num_scheduled_tokens, + { + spec_req_id: len(toks) + for spec_req_id, toks in scheduled_spec_tokens.items() + }, + ) + if new_token_ids: + req_state.output_token_ids.extend( + new_token_ids[-num_new_tokens:] + ) elif num_output_tokens < len(req_state.output_token_ids): # Some output tokens were discarded due to a sync-KV-load # failure, or output_token_ids was inflated by the optimistic @@ -1331,7 +2416,7 @@ class GPUModelRunner( reqs_to_add.append(req_state) # Track resumed requests for ngram_gpu full tensor copy - if is_ngram_gpu: + if has_ngram_gpu_buffers: ngram_gpu_new_reqs.append(req_state) continue @@ -1339,6 +2424,23 @@ class GPUModelRunner( self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens if new_block_ids is not None: self.input_batch.block_table.append_row(new_block_ids, req_index) + if trace_cow_worker: + self._trace_xpu_cow_worker_event( + stage="after_persistent_batch_update", + scheduler_output=scheduler_output, + req_id=req_id, + req_state=req_state, + req_index=req_index, + extra={ + "scheduled_cached_index": int(i), + "scheduler_num_computed_tokens": int( + num_computed_tokens + ), + "new_block_ids_lens": None + if new_block_ids is None + else [len(ids) for ids in new_block_ids], + }, + ) # For the last rank, we don't need to update the token_ids_cpu # because the sampled tokens are already cached. @@ -1352,7 +2454,21 @@ class GPUModelRunner( self.input_batch.num_tokens_no_spec[req_index] = end_token_index # Add spec_token_ids to token_ids_cpu. - self.input_batch.update_req_spec_token_ids(req_state, scheduled_spec_tokens) + spec_update_trace = self.input_batch.update_req_spec_token_ids( + req_state, scheduled_spec_tokens, trace=trace_cow_worker + ) + if trace_cow_worker: + self._trace_xpu_cow_worker_event( + stage="after_spec_token_update", + scheduler_output=scheduler_output, + req_id=req_id, + req_state=req_state, + req_index=req_index, + extra={ + "scheduled_cached_index": int(i), + "spec_update": spec_update_trace, + }, + ) # Restore scheduler-side draft count after ngram trimming. if original_num_spec_per_req: orig = original_num_spec_per_req.get(req_id, 0) @@ -1363,7 +2479,19 @@ class GPUModelRunner( # The smaller empty indices are filled first. for request in reqs_to_add: self.input_batch.add_request(request) - self.input_batch.update_req_spec_token_ids(request, scheduled_spec_tokens) + req_index = self.input_batch.req_id_to_index.get(request.req_id) + spec_update_trace = self.input_batch.update_req_spec_token_ids( + request, scheduled_spec_tokens, trace=trace_cow_worker + ) + if trace_cow_worker: + self._trace_xpu_cow_worker_event( + stage="after_new_req_add_spec_update", + scheduler_output=scheduler_output, + req_id=request.req_id, + req_state=request, + req_index=req_index, + extra={"spec_update": spec_update_trace}, + ) # Condense the batched states if there are gaps left by removed requests self.input_batch.condense() @@ -1373,7 +2501,7 @@ class GPUModelRunner( self.input_batch.refresh_metadata() # Incrementally update ngram_gpu tensors after batch is stable - if is_ngram_gpu: + if has_ngram_gpu_buffers: update_ngram_gpu_tensors_incremental( self.input_batch, self.token_ids_gpu_tensor, @@ -1410,7 +2538,7 @@ class GPUModelRunner( self.input_batch.num_computed_tokens_cpu[cur_req_index] -= ( correction ) - if is_ngram_gpu and correction > 0: + if has_ngram_gpu_buffers and correction > 0: self.input_batch.num_tokens_no_spec[cur_req_index] -= correction self.num_tokens_no_spec_gpu[cur_req_index] -= correction @@ -2010,6 +3138,54 @@ class GPUModelRunner( self.query_start_loc.gpu[: num_reqs + 1], self.positions[:total_num_scheduled_tokens], ) + if self._xpu_cow_worker_trace_enabled(): + self._trace_xpu_cow_worker_event( + stage="after_prepare_positions", + scheduler_output=scheduler_output, + extra={ + "num_reqs": int(num_reqs), + "total_num_scheduled_tokens": int( + total_num_scheduled_tokens + ), + "num_scheduled_tokens_head": self._cow_worker_array_head( + num_scheduled_tokens[:num_reqs] + ), + "req_indices_head": self._cow_worker_array_head( + req_indices[:total_num_scheduled_tokens] + ), + "query_pos_head": self._cow_worker_array_head( + self.query_pos.np[:total_num_scheduled_tokens] + ), + "num_computed_tokens_gpu_head": ( + self._cow_worker_tensor_head( + self.num_computed_tokens[:num_reqs] + ) + ), + "positions_head": self._cow_worker_tensor_head( + self.positions[:total_num_scheduled_tokens] + ), + "seq_lens_head": self._cow_worker_tensor_head( + self.seq_lens[:num_reqs] + ), + "prev_positions_head": self._cow_worker_array_head( + self.prev_positions.np[:num_reqs] + ), + "prev_num_draft_tokens_head": ( + self._cow_worker_array_head( + self.prev_num_draft_tokens.np[:num_reqs] + ) + ), + "request_windows": ( + self._cow_worker_prepare_request_windows( + scheduler_output=scheduler_output, + num_reqs=num_reqs, + num_scheduled_tokens=num_scheduled_tokens, + cu_num_tokens=cu_num_tokens, + req_indices=req_indices, + ) + ), + }, + ) # Copy the tensors to the GPU. self._prepare_input_ids( @@ -2095,6 +3271,976 @@ class GPUModelRunner( spec_decode_metadata, ) + def _xpu_cow_worker_trace_enabled(self) -> bool: + if not self.xpu_cow_worker_trace_file: + return False + if ( + self.xpu_cow_worker_trace_max_lines > 0 + and self.xpu_cow_worker_trace_lines + >= self.xpu_cow_worker_trace_max_lines + ): + return False + if self.xpu_cow_worker_trace_rank not in ("", "*", "all"): + if str(self.tp_rank) != self.xpu_cow_worker_trace_rank: + return False + return True + + def _cow_worker_tensor_head( + self, tensor: torch.Tensor | None, limit: int = 64 + ) -> dict[str, Any] | None: + if tensor is None: + return None + try: + detached = tensor.detach() + flat = detached.reshape(-1) + return { + "shape": list(detached.shape), + "dtype": str(detached.dtype), + "device": str(detached.device), + "head": flat[:limit].to("cpu").tolist(), + } + except Exception as exc: + return {"error": repr(exc)} + + def _cow_worker_tensor_slice( + self, + tensor: torch.Tensor | None, + start: int, + end: int, + limit: int = 64, + ) -> dict[str, Any] | None: + if tensor is None: + return None + try: + start = max(0, int(start)) + end = max(start, int(end)) + captured_end = min(end, start + int(limit)) + detached = tensor[start:captured_end].detach() + flat = detached.reshape(-1) + return { + "start": start, + "end": end, + "captured_end": captured_end, + "shape": list(detached.shape), + "dtype": str(detached.dtype), + "device": str(detached.device), + "head": flat.to("cpu").tolist(), + } + except Exception as exc: + return {"error": repr(exc), "start": start, "end": end} + + def _cow_worker_array_head( + self, array: Any, limit: int = 64 + ) -> dict[str, Any]: + try: + np_array = np.asarray(array) + flat = np_array.reshape(-1) + return { + "shape": list(np_array.shape), + "dtype": str(np_array.dtype), + "head": flat[:limit].tolist(), + } + except Exception as exc: + return {"error": repr(exc)} + + def _cow_worker_request_record( + self, req_state: CachedRequestState | None + ) -> dict[str, Any] | None: + if req_state is None: + return None + try: + return { + "req_id": req_state.req_id, + "num_prompt_tokens": int(req_state.num_prompt_tokens), + "num_output_tokens": len(req_state.output_token_ids), + "num_tokens": int(req_state.num_tokens), + "num_computed_tokens": int(req_state.num_computed_tokens), + "prev_num_draft_len": int(req_state.prev_num_draft_len), + "block_ids_lens": [ + len(block_ids) for block_ids in req_state.block_ids + ], + "block_ids_head": [ + list(block_ids[:8]) for block_ids in req_state.block_ids + ], + "output_token_ids_tail": req_state.output_token_ids[-16:], + } + except Exception as exc: + return {"error": repr(exc)} + + def _cow_worker_input_row_record( + self, req_index: int | None, token_window: int = 16 + ) -> dict[str, Any] | None: + if req_index is None: + return None + try: + record: dict[str, Any] = { + "req_index": int(req_index), + "req_id": self.input_batch.req_ids[req_index], + "num_prompt_tokens": int( + self.input_batch.num_prompt_tokens[req_index] + ), + "num_tokens_no_spec": int( + self.input_batch.num_tokens_no_spec[req_index] + ), + "num_computed_tokens_cpu": int( + self.input_batch.num_computed_tokens_cpu[req_index] + ), + "num_accepted_tokens_cpu": int( + self.input_batch.num_accepted_tokens_cpu[req_index] + ), + "spec_token_ids_head": list( + self.input_batch.spec_token_ids[req_index][:16] + ), + } + center = record["num_tokens_no_spec"] + window_start = max(0, int(center) - token_window) + window_end = min(self.max_model_len, int(center) + token_window) + record["token_ids_window_start"] = int(window_start) + record["token_ids_window"] = self.input_batch.token_ids_cpu[ + req_index, window_start:window_end + ].tolist() + + block_rows: list[dict[str, Any]] = [] + for group_index, table in enumerate( + self.input_batch.block_table.block_tables + ): + num_blocks = int(table.num_blocks_per_row[req_index]) + block_rows.append( + { + "group": int(group_index), + "num_blocks": num_blocks, + "head": table.block_table.np[ + req_index, : min(num_blocks, 16) + ].tolist(), + } + ) + record["block_table_rows"] = block_rows + return record + except Exception as exc: + return {"error": repr(exc), "req_index": req_index} + + def _cow_worker_prepare_request_windows( + self, + *, + scheduler_output: "SchedulerOutput", + num_reqs: int, + num_scheduled_tokens: np.ndarray, + cu_num_tokens: np.ndarray, + req_indices: np.ndarray, + max_reqs: int = 16, + max_tokens_per_req: int = 32, + token_window: int = 8, + ) -> list[dict[str, Any]]: + windows: list[dict[str, Any]] = [] + scheduler_spec_tokens = scheduler_output.scheduled_spec_decode_tokens + try: + capped_reqs = min(int(num_reqs), int(max_reqs)) + for req_index in range(capped_reqs): + req_id = self.input_batch.req_ids[req_index] + scheduled_count = int(num_scheduled_tokens[req_index]) + batch_start = ( + 0 if req_index == 0 else int(cu_num_tokens[req_index - 1]) + ) + batch_end = int(cu_num_tokens[req_index]) + captured_end = min(batch_end, batch_start + max_tokens_per_req) + + scheduled_spec = scheduler_spec_tokens.get(req_id, ()) + num_tokens_no_spec = int( + self.input_batch.num_tokens_no_spec[req_index] + ) + spec_write_start = ( + num_tokens_no_spec if scheduled_spec else None + ) + spec_write_end = ( + num_tokens_no_spec + len(scheduled_spec) + if scheduled_spec + else None + ) + + token_window_start = max( + 0, + num_tokens_no_spec - int(token_window), + ) + token_window_end = min( + self.max_model_len, + num_tokens_no_spec + + max(len(scheduled_spec), 1) + + int(token_window), + ) + + positions_cpu: list[int] = [] + try: + positions_cpu = [ + int(pos) + for pos in self.positions[ + batch_start:captured_end + ].detach().to("cpu").tolist() + ] + except Exception: + positions_cpu = [] + + token_ids_at_positions: list[dict[str, Any]] = [] + for pos in positions_cpu: + item: dict[str, Any] = {"position": pos} + if 0 <= pos < self.max_model_len: + item["token_id"] = int( + self.input_batch.token_ids_cpu[req_index, pos] + ) + else: + item["token_id"] = None + token_ids_at_positions.append(item) + + block_and_slot_windows: list[dict[str, Any]] = [] + for group_index, table in enumerate( + self.input_batch.block_table.block_tables + ): + num_blocks = int(table.num_blocks_per_row[req_index]) + block_tail_start = max(0, num_blocks - 8) + block_and_slot_windows.append( + { + "group": int(group_index), + "block_size": int(table.block_size), + "num_blocks": num_blocks, + "block_tail_start": int(block_tail_start), + "block_tail": table.block_table.np[ + req_index, block_tail_start:num_blocks + ].tolist(), + "slot_mapping": self._cow_worker_tensor_slice( + table.slot_mapping.gpu, + batch_start, + batch_end, + max_tokens_per_req, + ), + } + ) + + windows.append( + { + "req_index": int(req_index), + "req_id": req_id, + "batch_token_offset_start": int(batch_start), + "batch_token_offset_end": int(batch_end), + "batch_token_offset_captured_end": int(captured_end), + "num_scheduled_tokens": scheduled_count, + "scheduled_spec_len": int(len(scheduled_spec)), + "scheduled_spec_head": list(scheduled_spec[:16]), + "num_prompt_tokens": int( + self.input_batch.num_prompt_tokens[req_index] + ), + "num_computed_tokens_cpu": int( + self.input_batch + .num_computed_tokens_cpu[req_index] + ), + "num_tokens_no_spec": num_tokens_no_spec, + "prev_num_draft_len": int( + self.requests[req_id].prev_num_draft_len + ) if req_id in self.requests else None, + "spec_write_start": spec_write_start, + "spec_write_end": spec_write_end, + "token_window_start": int(token_window_start), + "token_window": self.input_batch.token_ids_cpu[ + req_index, token_window_start:token_window_end + ].tolist(), + "req_indices": np.asarray( + req_indices[batch_start:captured_end] + ).tolist(), + "query_pos": np.asarray( + self.query_pos.np[batch_start:captured_end] + ).tolist(), + "positions": self._cow_worker_tensor_slice( + self.positions, + batch_start, + batch_end, + max_tokens_per_req, + ), + "token_ids_at_positions": token_ids_at_positions, + "block_and_slot_windows": block_and_slot_windows, + } + ) + except Exception as exc: + windows.append({"error": repr(exc)}) + return windows + + def _trace_xpu_cow_worker_event( + self, + *, + stage: str, + scheduler_output: "SchedulerOutput", + req_id: str | None = None, + req_state: CachedRequestState | None = None, + req_index: int | None = None, + extra: dict[str, Any] | None = None, + ) -> None: + if not self._xpu_cow_worker_trace_enabled(): + return + + scheduler_spec_tokens = scheduler_output.scheduled_spec_decode_tokens + record: dict[str, Any] = { + "ts": time.time(), + "pid": os.getpid(), + "tp_rank": self.tp_rank, + "dcp_rank": self.dcp_rank, + "device": str(self.device), + "stage": stage, + "req_id": req_id, + "request": self._cow_worker_request_record(req_state), + "input_row": self._cow_worker_input_row_record(req_index), + "scheduler": { + "total_num_scheduled_tokens": int( + scheduler_output.total_num_scheduled_tokens + ), + "num_scheduled_tokens": { + sched_req_id: int(count) + for sched_req_id, count in + scheduler_output.num_scheduled_tokens.items() + }, + "scheduled_spec_decode_tokens": { + sched_req_id: { + "len": len(tokens), + "head": list(tokens[:16]), + } + for sched_req_id, tokens in scheduler_spec_tokens.items() + }, + }, + } + if extra: + record["extra"] = extra + num_reqs = extra.get("num_reqs") + if isinstance(num_reqs, int): + record["input_batch_rows"] = [ + self._cow_worker_input_row_record(i) + for i in range(min(num_reqs, 16)) + ] + + try: + with open(self.xpu_cow_worker_trace_file, "a", + encoding="utf-8") as f: + f.write(json.dumps(record, separators=(",", ":")) + "\n") + self.xpu_cow_worker_trace_lines += 1 + except Exception: + logger.exception("Failed to write XPU COW worker trace record") + + def _trace_xpu_model_inputs( + self, + *, + scheduler_output: "SchedulerOutput", + num_scheduled_tokens_np: np.ndarray, + batch_desc: BatchDescriptor, + cudagraph_mode: CUDAGraphMode, + input_ids: torch.Tensor | None, + positions: torch.Tensor | None, + logits_indices: torch.Tensor, + slot_mappings_by_group: dict[int, torch.Tensor] | None, + num_tokens_unpadded: int, + num_tokens_padded: int, + num_reqs: int, + use_spec_decode: bool, + skip_compiled: bool, + ) -> None: + if not self.xpu_model_input_trace_file: + return + if ( + self.xpu_model_input_trace_max_lines > 0 + and self.xpu_model_input_trace_lines + >= self.xpu_model_input_trace_max_lines + ): + return + if self.xpu_model_input_trace_rank not in ("", "*", "all"): + if str(self.tp_rank) != self.xpu_model_input_trace_rank: + return + + def tensor_head( + tensor: torch.Tensor | None, limit: int = 64 + ) -> dict[str, Any] | None: + if tensor is None: + return None + try: + detached = tensor.detach() + flat = detached.reshape(-1) + return { + "shape": list(detached.shape), + "dtype": str(detached.dtype), + "device": str(detached.device), + "head": flat[:limit].to("cpu").tolist(), + } + except Exception as exc: + return {"error": repr(exc)} + + def array_head(array: np.ndarray, limit: int = 64) -> dict[str, Any]: + flat = array.reshape(-1) + return { + "shape": list(array.shape), + "dtype": str(array.dtype), + "head": flat[:limit].tolist(), + } + + def block_table_head( + table: Any, + group_index: int, + num_reqs: int, + limit: int = 64, + ) -> dict[str, Any]: + record: dict[str, Any] = { + "group": group_index, + "type": type(table).__name__, + } + if torch.is_tensor(table): + tensor_record = tensor_head(table[:num_reqs], limit=limit) + if tensor_record is not None: + record["tensor"] = tensor_record + return record + + try: + num_blocks = getattr(table, "num_blocks_per_row", None) + if num_blocks is not None: + record["num_blocks_per_row"] = array_head( + np.asarray(num_blocks[:num_reqs]), limit=limit + ) + except Exception as exc: + record["num_blocks_per_row_error"] = repr(exc) + + accessors = ( + ("cpu", "get_cpu_tensor", False), + ("numpy", "get_numpy_array", False), + ("device", "get_device_tensor", True), + ("gpu", "get_gpu_tensor", True), + ) + for label, accessor_name, needs_num_reqs in accessors: + accessor = getattr(table, accessor_name, None) + if accessor is None: + continue + try: + value = accessor(num_reqs) if needs_num_reqs else accessor() + if torch.is_tensor(value): + record[label] = tensor_head(value[:num_reqs], limit=limit) + else: + record[label] = array_head( + np.asarray(value[:num_reqs]), limit=limit + ) + except Exception as exc: + record[f"{label}_error"] = repr(exc) + return record + + def call_bool(obj: Any, name: str) -> bool | str | None: + method = getattr(obj, name, None) + if method is None: + return None + try: + return bool(method()) + except Exception as exc: + return f"error:{exc!r}" + + def spec_config_record() -> dict[str, Any] | None: + spec_config = self.speculative_config + if spec_config is None: + return None + return { + "method": getattr(spec_config, "method", None), + "num_speculative_tokens": getattr( + spec_config, "num_speculative_tokens", None + ), + "draft_model_config_present": getattr( + spec_config, "draft_model_config", None + ) + is not None, + "uses_draft_model": call_bool(spec_config, "uses_draft_model"), + "use_ngram_gpu": call_bool(spec_config, "use_ngram_gpu"), + "use_eagle": call_bool(spec_config, "use_eagle"), + "use_dflash": call_bool(spec_config, "use_dflash"), + "uses_extract_hidden_states": call_bool( + spec_config, "uses_extract_hidden_states" + ), + "enforce_eager": getattr(spec_config, "enforce_eager", None), + } + + def cache_config_record() -> dict[str, Any]: + record: dict[str, Any] = { + "cache_dtype": getattr(self.cache_config, "cache_dtype", None), + "kv_cache_dtype": str(self.kv_cache_dtype), + "configured_block_size": getattr( + self.cache_config, "block_size", None + ), + } + kv_cache_config = getattr(self, "kv_cache_config", None) + if kv_cache_config is None: + return record + groups = getattr(kv_cache_config, "kv_cache_groups", []) or [] + record["num_kv_cache_groups"] = len(groups) + group_records: list[dict[str, Any]] = [] + for group_index, group in enumerate(groups): + group_record: dict[str, Any] = {"group": group_index} + spec = getattr(group, "kv_cache_spec", None) + if spec is not None: + group_record["spec_type"] = type(spec).__name__ + for attr in ( + "block_size", + "num_speculative_blocks", + "sliding_window", + ): + if hasattr(spec, attr): + try: + group_record[attr] = getattr(spec, attr) + except Exception as exc: + group_record[f"{attr}_error"] = repr(exc) + group_records.append(group_record) + record["groups"] = group_records + return record + + def request_state_records(req_ids: list[str]) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + for req_id in req_ids: + state = self.requests.get(req_id) + if state is None: + records.append({"req_id": req_id, "missing": True}) + continue + records.append( + { + "req_id": req_id, + "num_prompt_tokens": int(state.num_prompt_tokens), + "num_output_tokens": len(state.output_token_ids), + "num_tokens": int(state.num_tokens), + "num_computed_tokens": int(state.num_computed_tokens), + "prev_num_draft_len": int(state.prev_num_draft_len), + "block_ids_lens": [ + len(block_ids) for block_ids in state.block_ids + ], + "block_ids_head": [ + list(block_ids[:8]) for block_ids in state.block_ids + ], + "output_token_ids_head": state.output_token_ids[-16:], + } + ) + return records + + block_table_records: list[dict[str, Any]] = [] + try: + for group_index in range(len(self.kv_cache_config.kv_cache_groups)): + table = self.input_batch.block_table[group_index] + block_table_records.append( + block_table_head(table, group_index, num_reqs, limit=64) + ) + except Exception as exc: + block_table_records.append({"error": repr(exc)}) + + slot_mapping_records: list[dict[str, Any]] = [] + if slot_mappings_by_group is not None: + for group_index, mapping in slot_mappings_by_group.items(): + record = tensor_head(mapping, limit=64) + if record is not None: + record["group"] = group_index + slot_mapping_records.append(record) + + req_ids = list(self.input_batch.req_ids[:num_reqs]) + spec_tokens = scheduler_output.scheduled_spec_decode_tokens + record = { + "ts": time.time(), + "pid": os.getpid(), + "tp_rank": self.tp_rank, + "dcp_rank": self.dcp_rank, + "device": str(self.device), + "cudagraph_mode": str(cudagraph_mode), + "config": { + "max_model_len": int(self.max_model_len), + "max_num_tokens": int(self.max_num_tokens), + "max_num_reqs": int(self.max_num_reqs), + "num_spec_tokens": int(self.num_spec_tokens), + "use_async_scheduling": bool(self.use_async_scheduling), + "use_async_spec_decode": bool(self.use_async_spec_decode), + "speculative": spec_config_record(), + "cache": cache_config_record(), + }, + "batch_desc": { + "num_tokens": int(batch_desc.num_tokens), + "num_reqs": None + if batch_desc.num_reqs is None + else int(batch_desc.num_reqs), + "uniform": bool(batch_desc.uniform), + }, + "scheduler": { + "total_num_scheduled_tokens": int( + scheduler_output.total_num_scheduled_tokens + ), + "num_scheduled_tokens": { + req_id: int(count) + for req_id, count in + scheduler_output.num_scheduled_tokens.items() + }, + "scheduled_spec_decode_tokens": { + req_id: { + "len": len(tokens), + "head": list(tokens[:16]), + } + for req_id, tokens in spec_tokens.items() + }, + }, + "input_batch": { + "req_ids": req_ids, + "num_reqs": int(num_reqs), + "num_tokens_unpadded": int(num_tokens_unpadded), + "num_tokens_padded": int(num_tokens_padded), + "num_scheduled_tokens_np": array_head( + num_scheduled_tokens_np[:num_reqs] + ), + "num_tokens_no_spec": array_head( + self.input_batch.num_tokens_no_spec[:num_reqs] + ), + "num_prompt_tokens_cpu": array_head( + self.input_batch.num_prompt_tokens[:num_reqs] + ), + "num_computed_tokens_cpu": array_head( + self.input_batch.num_computed_tokens_cpu[:num_reqs] + ), + "num_accepted_tokens_cpu": array_head( + self.input_batch.num_accepted_tokens_cpu[:num_reqs] + ), + "spec_token_ids": [ + list(tokens[:16]) + for tokens in self.input_batch.spec_token_ids[:num_reqs] + ], + "request_states": request_state_records(req_ids), + "input_ids": tensor_head(input_ids), + "positions": tensor_head(positions), + "logits_indices": tensor_head(logits_indices), + "use_spec_decode": bool(use_spec_decode), + "skip_compiled": bool(skip_compiled), + }, + "attn": { + "block_tables": block_table_records, + "slot_mappings": slot_mapping_records, + }, + } + try: + with open(self.xpu_model_input_trace_file, "a", encoding="utf-8") as f: + f.write(json.dumps(record, separators=(",", ":")) + "\n") + self.xpu_model_input_trace_lines += 1 + except Exception: + logger.exception("Failed to write XPU model input trace") + + def _xpu_replay_microscope_req_ids(self, num_reqs: int | None) -> list[str]: + try: + count = self.input_batch.num_reqs if num_reqs is None else num_reqs + return list(self.input_batch.req_ids[:count]) + except Exception: + return [] + + def _xpu_replay_microscope_matches( + self, req_ids: list[str] + ) -> tuple[bool, list[str]]: + if not self.xpu_replay_microscope_file: + return False, [] + if ( + self.xpu_replay_microscope_max_lines > 0 + and self.xpu_replay_microscope_lines + >= self.xpu_replay_microscope_max_lines + ): + return False, [] + if self.xpu_replay_microscope_rank not in ("", "*", "all"): + if str(self.tp_rank) != self.xpu_replay_microscope_rank: + return False, [] + req_regex = self.xpu_replay_microscope_req_regex + if not req_regex: + return True, req_ids + try: + matched = [req_id for req_id in req_ids if re.search(req_regex, req_id)] + except re.error as exc: + matched = [] + logger.warning("Invalid VLLM_XPU_REPLAY_MICROSCOPE_REQ_REGEX: %r", exc) + return bool(matched), matched + + @staticmethod + def _xpu_replay_json_number(value: Any) -> int | float | str | None: + try: + if value is None: + return None + if isinstance(value, (np.integer,)): + return int(value) + if isinstance(value, (np.floating,)): + value = float(value) + if isinstance(value, bool): + return bool(value) + if isinstance(value, int): + return int(value) + if isinstance(value, float): + if np.isfinite(value): + return value + if np.isnan(value): + return "nan" + if value == float("inf"): + return "inf" + if value == -float("inf"): + return "-inf" + return repr(value) + except Exception: + return repr(value) + return value + + def _xpu_replay_tensor_record( + self, + tensor: torch.Tensor | None, + *, + include_head: bool = True, + include_topk: bool = False, + ) -> dict[str, Any] | None: + if tensor is None: + return None + limit = max(0, int(self.xpu_replay_microscope_tensor_limit)) + topk = max(0, int(self.xpu_replay_microscope_topk)) + try: + with torch.no_grad(): + t = tensor.detach() + record: dict[str, Any] = { + "shape": list(t.shape), + "dtype": str(t.dtype), + "device": str(t.device), + "numel": int(t.numel()), + } + if t.numel() == 0: + record.update( + {"finite": 0, "nan": 0, "posinf": 0, "neginf": 0} + ) + return record + + stats_t = t.float() if not t.is_floating_point() else t + finite_mask = torch.isfinite(stats_t) + nan_mask = torch.isnan(stats_t) + posinf_mask = stats_t == float("inf") + neginf_mask = stats_t == -float("inf") + finite_count = int(finite_mask.sum().item()) + record.update( + { + "finite": finite_count, + "nan": int(nan_mask.sum().item()), + "posinf": int(posinf_mask.sum().item()), + "neginf": int(neginf_mask.sum().item()), + } + ) + if finite_count: + finite_vals = stats_t[finite_mask].float() + record.update( + { + "min": self._xpu_replay_json_number( + finite_vals.min().item() + ), + "max": self._xpu_replay_json_number( + finite_vals.max().item() + ), + "mean": self._xpu_replay_json_number( + finite_vals.mean().item() + ), + "sum": self._xpu_replay_json_number( + finite_vals.sum().item() + ), + "l2": self._xpu_replay_json_number( + torch.linalg.vector_norm(finite_vals).item() + ), + } + ) + if include_head and limit > 0: + head = t.reshape(-1)[:limit].to("cpu").tolist() + record["head"] = [ + self._xpu_replay_json_number(value) for value in head + ] + if include_topk and topk > 0 and t.dim() >= 2 and t.shape[-1] > 0: + rows = min(int(t.shape[0]), max(1, limit)) + k = min(topk, int(t.shape[-1])) + row_records: list[dict[str, Any]] = [] + for row_index in range(rows): + row = stats_t[row_index].reshape(-1) + values, indices = torch.topk(row, k=k, dim=-1) + row_records.append( + { + "row": int(row_index), + "token_ids": [ + int(value) + for value in indices.to("cpu").tolist() + ], + "values": [ + self._xpu_replay_json_number(value) + for value in values.to("cpu").tolist() + ], + } + ) + record["topk"] = row_records + return record + except Exception as exc: + return {"error": repr(exc)} + + def _xpu_replay_request_records( + self, + req_ids: list[str], + matched_req_ids: list[str], + ) -> list[dict[str, Any]]: + matched = set(matched_req_ids) if matched_req_ids else set(req_ids) + records: list[dict[str, Any]] = [] + for req_index, req_id in enumerate(req_ids): + if req_id not in matched: + continue + record: dict[str, Any] = {"req_index": int(req_index), "req_id": req_id} + try: + record["num_prompt_tokens_cpu"] = int( + self.input_batch.num_prompt_tokens[req_index] + ) + record["num_computed_tokens_cpu"] = int( + self.input_batch.num_computed_tokens_cpu[req_index] + ) + record["num_tokens_no_spec"] = int( + self.input_batch.num_tokens_no_spec[req_index] + ) + total_tokens = record["num_tokens_no_spec"] + start = max(0, total_tokens - 24) + stop = min( + self.input_batch.token_ids_cpu.shape[1], + max(total_tokens + 4, start), + ) + record["token_ids_cpu_window_start"] = int(start) + record["token_ids_cpu_window"] = self.input_batch.token_ids_cpu[ + req_index, start:stop + ].tolist() + except Exception as exc: + record["input_batch_error"] = repr(exc) + try: + state = self.requests.get(req_id) + if state is None: + record["request_state_missing"] = True + else: + record["state_num_tokens"] = int(state.num_tokens) + record["state_num_computed_tokens"] = int( + state.num_computed_tokens + ) + record["state_output_token_ids_tail"] = list( + state.output_token_ids[-24:] + ) + record["state_prev_num_draft_len"] = int( + state.prev_num_draft_len + ) + record["state_block_ids_tail"] = [ + list(block_ids[-8:]) for block_ids in state.block_ids + ] + except Exception as exc: + record["request_state_error"] = repr(exc) + records.append(record) + return records + + def _trace_xpu_replay_microscope( + self, + *, + stage: str, + scheduler_output: "SchedulerOutput | None" = None, + num_reqs: int | None = None, + batch_desc: BatchDescriptor | None = None, + cudagraph_mode: CUDAGraphMode | None = None, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + logits_indices: torch.Tensor | None = None, + slot_mappings_by_group: dict[int, torch.Tensor] | None = None, + hidden_states: torch.Tensor | None = None, + sample_hidden_states: torch.Tensor | None = None, + logits: torch.Tensor | None = None, + sampler_output: SamplerOutput | None = None, + extra: dict[str, Any] | None = None, + ) -> None: + req_ids = self._xpu_replay_microscope_req_ids(num_reqs) + enabled, matched_req_ids = self._xpu_replay_microscope_matches(req_ids) + if not enabled: + return + + record: dict[str, Any] = { + "ts": time.time(), + "pid": os.getpid(), + "tp_rank": self.tp_rank, + "dcp_rank": self.dcp_rank, + "device": str(self.device), + "stage": stage, + "req_ids": req_ids, + "matched_req_ids": matched_req_ids, + "requests": self._xpu_replay_request_records( + req_ids, matched_req_ids + ), + } + if scheduler_output is not None: + try: + record["scheduler"] = { + "total_num_scheduled_tokens": int( + scheduler_output.total_num_scheduled_tokens + ), + "num_scheduled_tokens": { + req_id: int(count) + for req_id, count in + scheduler_output.num_scheduled_tokens.items() + }, + "scheduled_spec_decode_tokens": { + req_id: { + "len": len(tokens), + "head": list(tokens[:16]), + } + for req_id, tokens in + scheduler_output.scheduled_spec_decode_tokens.items() + }, + } + except Exception as exc: + record["scheduler_error"] = repr(exc) + if batch_desc is not None: + record["batch_desc"] = { + "num_tokens": int(batch_desc.num_tokens), + "num_reqs": None + if batch_desc.num_reqs is None + else int(batch_desc.num_reqs), + "uniform": bool(batch_desc.uniform), + } + if cudagraph_mode is not None: + record["cudagraph_mode"] = str(cudagraph_mode) + if extra: + record["extra"] = extra + + tensor_records: dict[str, Any] = {} + for name, tensor in ( + ("input_ids", input_ids), + ("positions", positions), + ("logits_indices", logits_indices), + ("hidden_states", hidden_states), + ("sample_hidden_states", sample_hidden_states), + ("logits", logits), + ): + if tensor is not None: + tensor_records[name] = self._xpu_replay_tensor_record( + tensor, + include_topk=(name == "logits"), + ) + if sampler_output is not None: + tensor_records["sampled_token_ids"] = ( + self._xpu_replay_tensor_record( + sampler_output.sampled_token_ids, + include_topk=False, + ) + ) + if sampler_output.logprobs_tensors is not None: + try: + tensor_records["logprobs_tensors"] = str( + type(sampler_output.logprobs_tensors) + ) + except Exception: + pass + if slot_mappings_by_group is not None: + slot_records: dict[str, Any] = {} + for group_index, mapping in slot_mappings_by_group.items(): + slot_records[str(group_index)] = self._xpu_replay_tensor_record( + mapping, + include_topk=False, + ) + tensor_records["slot_mappings_by_group"] = slot_records + if tensor_records: + record["tensors"] = tensor_records + + try: + with open(self.xpu_replay_microscope_file, "a", encoding="utf-8") as f: + f.write( + json.dumps(record, separators=(",", ":"), allow_nan=False) + + "\n" + ) + self.xpu_replay_microscope_lines += 1 + except Exception: + logger.exception("Failed to write XPU replay microscope trace") + def _build_attention_metadata( self, num_tokens: int, @@ -2310,8 +4456,9 @@ class GPUModelRunner( cm.slot_mapping = slot_mappings[kv_cache_gid] if self.speculative_config and spec_decode_common_attn_metadata is None: - if isinstance(self.drafter, (EagleProposer, DFlashProposer)): - if self.drafter.kv_cache_gid == kv_cache_gid: + drafter = getattr(self, "drafter", None) + if isinstance(drafter, (EagleProposer, DFlashProposer)): + if drafter.kv_cache_gid == kv_cache_gid: spec_decode_common_attn_metadata = cm else: spec_decode_common_attn_metadata = cm @@ -3355,13 +5502,77 @@ class GPUModelRunner( encoder_outputs = self._execute_mm_encoder(scheduler_output) model_kwargs.update({"encoder_outputs": encoder_outputs}) - return ( - input_ids, - inputs_embeds, - positions, - intermediate_tensors, - model_kwargs, - ec_connector_output, + return ( + input_ids, + inputs_embeds, + positions, + intermediate_tensors, + model_kwargs, + ec_connector_output, + ) + + def _xpu_local_argmax_debug(self, message: str) -> None: + if os.environ.get("VLLM_XPU_LOCAL_ARGMAX_DEBUG", "0") != "1": + return + count = getattr(self, "_xpu_local_argmax_debug_count", 0) + limit = int(os.environ.get("VLLM_XPU_LOCAL_ARGMAX_DEBUG_LIMIT", "24")) + if count >= limit: + return + self._xpu_local_argmax_debug_count = count + 1 + rank = -1 + if torch.distributed.is_available() and torch.distributed.is_initialized(): + rank = torch.distributed.get_rank() + logger.info("XPU local argmax rank=%s: %s", rank, message) + + def _sync_xpu_sampled_token_ids_across_tp( + self, + sampled_token_ids: torch.Tensor, + *, + env_name: str, + default: str, + timing_label: str, + strict_env_name: str, + ) -> torch.Tensor: + if os.environ.get(env_name, default) != "1": + return sampled_token_ids + if sampled_token_ids.device.type != "xpu": + return sampled_token_ids + if ( + not torch.distributed.is_available() + or not torch.distributed.is_initialized() + ): + return sampled_token_ids + try: + tp_group = get_tp_group() + if tp_group.world_size <= 1: + return sampled_token_ids + synced = sampled_token_ids.contiguous() + with timed_region(timing_label): + torch.distributed.broadcast( + synced, + src=tp_group.first_rank, + group=tp_group.device_group, + ) + self._xpu_local_argmax_debug( + f"synced sampled_token_ids from TP first rank via {env_name}" + ) + return synced + except Exception as exc: + if os.environ.get(strict_env_name, "0") == "1": + raise + self._xpu_local_argmax_debug(f"{env_name} skipped: {exc!r}") + return sampled_token_ids + + def _sync_xpu_local_argmax_sampled_token_ids( + self, + sampled_token_ids: torch.Tensor, + ) -> torch.Tensor: + return self._sync_xpu_sampled_token_ids_across_tp( + sampled_token_ids, + env_name="VLLM_XPU_LOCAL_ARGMAX_SYNC_TOKENS", + default="1", + timing_label="gpu_model_runner.local_argmax_sync_tokens", + strict_env_name="VLLM_XPU_LOCAL_ARGMAX_SYNC_STRICT", ) def _sample( @@ -3369,16 +5580,28 @@ class GPUModelRunner( logits: torch.Tensor | None, spec_decode_metadata: SpecDecodeMetadata | None, ) -> SamplerOutput: + precomputed_sampled_token_ids = getattr( + self, "_xpu_local_argmax_sampled_token_ids", None + ) + if logits is None and precomputed_sampled_token_ids is not None: + self._xpu_local_argmax_sampled_token_ids = None + self._xpu_local_argmax_debug("using precomputed sampled_token_ids") + return SamplerOutput( + sampled_token_ids=precomputed_sampled_token_ids, + logprobs_tensors=None, + ) + # Sample the next token and get logprobs if needed. sampling_metadata = self.input_batch.sampling_metadata # Update output token ids with tokens sampled in last step # if async scheduling and required by current sampling params. self.input_batch.update_async_output_token_ids() if spec_decode_metadata is None: - return self.sampler( - logits=logits, - sampling_metadata=sampling_metadata, - ) + with timed_region("gpu_model_runner.sampler"): + return self.sampler( + logits=logits, + sampling_metadata=sampling_metadata, + ) # Update spec_token_ids with real draft tokens from pre step only when # output_token_ids is needed (penalties or bad_words are in use). @@ -3386,14 +5609,74 @@ class GPUModelRunner( draft_token_ids_cpu, _ = self._get_draft_token_ids_cpu() self.input_batch.update_async_spec_token_ids(draft_token_ids_cpu) - sampler_output = self.rejection_sampler( - spec_decode_metadata, - None, # draft_probs - logits, - sampling_metadata, - ) + with timed_region("gpu_model_runner.rejection_sampler"): + sampler_output = self.rejection_sampler( + spec_decode_metadata, + None, # draft_probs + logits, + sampling_metadata, + ) return sampler_output + def _can_use_xpu_local_argmax( + self, + spec_decode_metadata: SpecDecodeMetadata | None, + ) -> bool: + def disabled(reason: str) -> bool: + self._xpu_local_argmax_debug(f"disabled: {reason}") + return False + + if os.environ.get("VLLM_XPU_LOCAL_ARGMAX_DECODE", "0") != "1": + return disabled("env flag is not enabled") + if spec_decode_metadata is not None: + return disabled("spec decode metadata is present") + if os.environ.get("VLLM_XPU_LOCAL_ARGMAX_ASSUME_SAFE", "0") == "1": + if not hasattr(self.model, "get_top_tokens"): + return disabled("model has no get_top_tokens method") + self._xpu_local_argmax_debug("enabled by assume-safe override") + return True + sampling_metadata = self.input_batch.sampling_metadata + if not sampling_metadata.all_greedy: + return disabled("not all requests are greedy") + if sampling_metadata.max_num_logprobs is not None: + return disabled("logprobs requested") + if sampling_metadata.logprob_token_ids: + return disabled("specific logprob token IDs requested") + if not sampling_metadata.no_penalties: + return disabled("sampling penalties are active") + if sampling_metadata.allowed_token_ids_mask is not None: + return disabled("allowed-token mask is active") + if sampling_metadata.bad_words_token_ids: + return disabled("bad-word filtering is active") + active_processors: list[str] = [] + for processor in sampling_metadata.logitsprocs.non_argmax_invariant: + name = processor.__class__.__name__ + if name == "LogitBiasLogitsProcessor": + if getattr(processor, "biases", None): + active_processors.append(name) + elif name == "MinTokensLogitsProcessor": + if getattr(processor, "min_toks", None): + active_processors.append(name) + elif name == "ThinkingTokenBudgetLogitsProcessor": + if getattr(processor, "is_enabled", False) and getattr( + processor, "_state", None + ): + active_processors.append(name) + else: + active_processors.append(name) + if active_processors: + return disabled( + "active non-argmax logits processors: " + + ",".join(active_processors) + ) + holder = sampling_metadata.thinking_budget_state_holder + if holder is not None and holder.has_tracked_requests(): + return disabled("thinking-budget tracking is active") + if not hasattr(self.model, "get_top_tokens"): + return disabled("model has no get_top_tokens method") + self._xpu_local_argmax_debug("enabled") + return True + def _bookkeeping_sync( self, scheduler_output: "SchedulerOutput", @@ -3477,17 +5760,45 @@ class GPUModelRunner( # the sampled tokens back, because there's no direct communication # between the first-stage worker and the last-stage worker. req_ids = self.input_batch.req_ids + xpu_filter_suppressed_bonus_cache = ( + os.environ.get( + "VLLM_XPU_SPEC_DECODE_FILTER_SUPPRESSED_BONUS_CACHE", "") + .strip() + .lower() + in ("1", "true", "yes", "on") + ) + scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens + trace_cow_worker = self._xpu_cow_worker_trace_enabled() for req_idx in range(num_sampled_tokens): if self.use_async_scheduling: sampled_ids = [-1] if req_idx not in invalid_req_indices_set else None else: sampled_ids = valid_sampled_token_ids[req_idx] - num_sampled_ids: int = len(sampled_ids) if sampled_ids else 0 - if not sampled_ids: continue + req_id = req_ids[req_idx] + sampled_ids_for_cache = sampled_ids + scheduled_spec_ids = scheduled_spec_tokens.get(req_id, ()) + if ( + xpu_filter_suppressed_bonus_cache + and scheduled_spec_ids + and len(sampled_ids) == len(scheduled_spec_ids) + 1 + ): + # Diagnostic mode: the scheduler receives the full sampler + # result for acceptance accounting, but it suppresses the + # verifier bonus token from user-visible output. Keep the + # worker-side token cache aligned with the emitted sequence so + # the next draft/proposer step does not advance one token ahead. + sampled_ids_for_cache = sampled_ids[:-1] + + num_sampled_ids: int = ( + len(sampled_ids_for_cache) if sampled_ids_for_cache else 0 + ) + if not sampled_ids_for_cache: + continue + start_idx = self.input_batch.num_tokens_no_spec[req_idx] end_idx = start_idx + num_sampled_ids assert end_idx <= self.max_model_len, ( @@ -3496,13 +5807,29 @@ class GPUModelRunner( f"{self.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_for_cache self.input_batch.is_token_ids[req_idx, start_idx:end_idx] = True self.input_batch.num_tokens_no_spec[req_idx] = end_idx - req_id = req_ids[req_idx] req_state = self.requests[req_id] - req_state.output_token_ids.extend(sampled_ids) + req_state.output_token_ids.extend(sampled_ids_for_cache) + if trace_cow_worker and sampled_ids_for_cache is not sampled_ids: + self._trace_xpu_cow_worker_event( + stage="after_sampled_cache_bonus_suppression", + scheduler_output=scheduler_output, + req_id=req_id, + req_state=req_state, + req_index=req_idx, + extra={ + "scheduled_spec_token_ids": list(scheduled_spec_ids), + "sampled_token_ids": list(sampled_ids), + "cached_token_ids": list(sampled_ids_for_cache), + "cache_start": int(start_idx), + "cache_end": int(end_idx), + }, + ) # Compute prompt logprobs if needed. prompt_logprobs_dict = self._get_prompt_logprobs_dict( @@ -3559,13 +5886,19 @@ class GPUModelRunner( Returns: Model output tensor """ - return self.model( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - **model_kwargs, - ) + with timed_region("gpu_model_runner.model_forward"): + return self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, + ) + + def _finish_xpu_timing_step(self, status: str = "ok") -> None: + if getattr(self, "_xpu_timing_step_active", False): + end_timing_step(status) + self._xpu_timing_step_active = False @staticmethod def _is_uniform_decode( @@ -3685,6 +6018,114 @@ class GPUModelRunner( # num_tokens_across_dp will no-longer be valid assert batch_descriptor.num_tokens == num_tokens_padded + if ( + os.environ.get("VLLM_XPU_DISABLE_PREFILL_CUDAGRAPH_REPLAY", "0") + == "1" + and force_uniform_decode is None + and not force_eager + and not uniform_decode + and cudagraph_mode == CUDAGraphMode.PIECEWISE + ): + # Qwen3.6/Quark W8A8 on XPU can produce different first-decode + # logits for identical fresh prompts when prefill executes through + # PIECEWISE replay. Keep decode replay intact and bypass only the + # non-uniform prefill graph boundary. + cudagraph_mode = CUDAGraphMode.NONE + elif ( + os.environ.get( + "VLLM_XPU_DISABLE_INITIAL_DECODE_CUDAGRAPH_REPLAY_STEPS", "0" + ) + not in ("", "0") + and uniform_decode + and cudagraph_mode == CUDAGraphMode.PIECEWISE + ): + try: + initial_decode_steps = int( + os.environ.get( + "VLLM_XPU_DISABLE_INITIAL_DECODE_CUDAGRAPH_REPLAY_STEPS", + "0", + ) + or "0" + ) + except Exception: + initial_decode_steps = 0 + if initial_decode_steps > 0: + try: + req_ids_cpu = self.input_batch.req_ids + num_reqs_cpu = min( + int(getattr(self.input_batch, "num_reqs", len(req_ids_cpu))), + len(req_ids_cpu), + len(self.input_batch.num_prompt_tokens), + len(self.input_batch.num_computed_tokens_cpu), + ) + prompt_tokens_cpu = self.input_batch.num_prompt_tokens + computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu + in_initial_decode_window = any( + req_ids_cpu[i] is not None + and 0 + <= int(computed_tokens_cpu[i]) - int(prompt_tokens_cpu[i]) + < initial_decode_steps + for i in range(num_reqs_cpu) + ) + except Exception: + in_initial_decode_window = False + if in_initial_decode_window: + # Bypass the unstable request boundary for the first few + # decode tokens, then re-enter PIECEWISE replay for the + # long steady-state decode tail. + cudagraph_mode = CUDAGraphMode.NONE + elif ( + os.environ.get("VLLM_XPU_DISABLE_DECODE_CUDAGRAPH_REPLAY", "0") == "1" + and uniform_decode + and cudagraph_mode == CUDAGraphMode.PIECEWISE + ): + # Diagnostic guard for XPU graph replay corruption: keep the same + # compiled runnable and padding decision, but bypass decode replay. + cudagraph_mode = CUDAGraphMode.NONE + elif ( + os.environ.get( + "VLLM_XPU_DISABLE_FIRST_DECODE_CUDAGRAPH_REPLAY", "0" + ) + == "1" + and uniform_decode + and cudagraph_mode == CUDAGraphMode.PIECEWISE + ): + try: + req_ids_cpu = self.input_batch.req_ids + num_reqs_cpu = min( + int(getattr(self.input_batch, "num_reqs", len(req_ids_cpu))), + len(req_ids_cpu), + len(self.input_batch.num_prompt_tokens), + len(self.input_batch.num_computed_tokens_cpu), + ) + prompt_tokens_cpu = self.input_batch.num_prompt_tokens + computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu + first_decode_replay = any( + req_ids_cpu[i] is not None + and int(computed_tokens_cpu[i]) == int(prompt_tokens_cpu[i]) + for i in range(num_reqs_cpu) + ) + except Exception: + first_decode_replay = False + if first_decode_replay: + # The Qwen3.6 XPU replay microscope showed identical eager + # prefill inputs diverging on the first PIECEWISE decode step. + # Bypass only that boundary to test whether replay is consuming + # stale/misaligned prefill state without paying the full + # no-decode-replay cost. + cudagraph_mode = CUDAGraphMode.NONE + elif uniform_decode and cudagraph_mode == CUDAGraphMode.PIECEWISE: + eager_every = int( + os.environ.get( + "VLLM_XPU_DECODE_CUDAGRAPH_REPLAY_EAGER_EVERY", "0" + ) + or "0" + ) + if eager_every > 0: + self.xpu_decode_cudagraph_replay_count += 1 + if self.xpu_decode_cudagraph_replay_count % eager_every == 0: + cudagraph_mode = CUDAGraphMode.NONE + cudagraph_stats = None if self.vllm_config.observability_config.cudagraph_metrics: cudagraph_stats = CUDAGraphStat( @@ -3832,6 +6273,8 @@ class GPUModelRunner( "State error: sample_tokens() must be called " "after execute_model() returns None." ) + if os.environ.get("VLLM_XPU_CUDAGRAPH_MARK_STEP_BEGIN", "0") == "1": + torch.compiler.cudagraph_mark_step_begin() if self.routed_experts_initialized: capturer = RoutedExpertsCapturer.get_instance() @@ -3865,10 +6308,28 @@ class GPUModelRunner( num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens with ( record_function_or_nullcontext("gpu_model_runner: preprocess"), + timed_region("gpu_model_runner.preprocess_total"), self.synchronize_input_prep(), ): # Update persistent batch states. deferred_state_corrections_fn = self._update_states(scheduler_output) + if os.environ.get("VLLM_XPU_TRACE_NGRAM_PP"): + logger.warning( + "XPU ngram post-update pp_rank=%s first=%s last=%s " + "pre_update_total=%s post_update_total=%s scheduled=%s " + "spec_lens=%s req_ids=%s", + get_pp_group().rank_in_group, + get_pp_group().is_first_rank, + get_pp_group().is_last_rank, + num_scheduled_tokens, + scheduler_output.total_num_scheduled_tokens, + scheduler_output.num_scheduled_tokens, + { + req_id: len(toks) + for req_id, toks in scheduler_output.scheduled_spec_decode_tokens.items() + }, + self.input_batch.req_ids, + ) if has_ec_transfer() and not get_ec_transfer().is_consumer: with self.maybe_get_ec_connector_output( @@ -4062,6 +6523,147 @@ class GPUModelRunner( has_encoder_input = ( self.model_config.is_encoder_decoder and num_encoder_reqs > 0 ) + skip_compiled = has_encoder_input or ( + os.environ.get("VLLM_XPU_SKIP_COMPILED_PREFILL", "0") == "1" + and num_tokens_padded > 1 + ) + self._trace_xpu_model_inputs( + scheduler_output=scheduler_output, + num_scheduled_tokens_np=num_scheduled_tokens_np, + batch_desc=batch_desc, + cudagraph_mode=cudagraph_mode, + input_ids=input_ids, + positions=positions, + logits_indices=logits_indices, + slot_mappings_by_group=slot_mappings_by_group, + num_tokens_unpadded=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded, + num_reqs=num_reqs, + use_spec_decode=use_spec_decode, + skip_compiled=skip_compiled, + ) + self._trace_xpu_replay_microscope( + stage="inputs", + scheduler_output=scheduler_output, + num_reqs=num_reqs, + batch_desc=batch_desc, + cudagraph_mode=cudagraph_mode, + input_ids=input_ids, + positions=positions, + logits_indices=logits_indices, + slot_mappings_by_group=slot_mappings_by_group, + extra={ + "num_tokens_unpadded": int(num_tokens_unpadded), + "num_tokens_padded": int(num_tokens_padded), + "max_num_scheduled_tokens": int(max_num_scheduled_tokens), + "use_spec_decode": bool(use_spec_decode), + "skip_compiled": bool(skip_compiled), + }, + ) + + scheduled_token_counts = [ + int(value) for value in num_scheduled_tokens_np[:num_reqs] + ] + scheduled_token_histogram: dict[str, int] = {} + for count in scheduled_token_counts: + key = str(count) + scheduled_token_histogram[key] = ( + scheduled_token_histogram.get(key, 0) + 1 + ) + + spec_tokens_by_req = scheduler_output.scheduled_spec_decode_tokens + scheduled_spec_lengths = [ + int(len(spec_tokens_by_req.get(req_id, ()))) + for req_id in req_ids[:num_reqs] + ] + scheduled_spec_histogram: dict[str, int] = {} + for count in scheduled_spec_lengths: + key = str(count) + scheduled_spec_histogram[key] = ( + scheduled_spec_histogram.get(key, 0) + 1 + ) + max_scheduled_spec_tokens = ( + max(scheduled_spec_lengths) if scheduled_spec_lengths else 0 + ) + + computed_token_counts = [ + int(value) + for value in self.input_batch.num_computed_tokens_cpu[:num_reqs] + ] + decode_req_count = sum(1 for count in computed_token_counts if count > 0) + prefill_req_count = int(num_reqs - decode_req_count) + is_pure_decode = prefill_req_count == 0 and num_reqs > 0 + decode_bucket = int(max_num_scheduled_tokens) if is_pure_decode else None + num_tokens_across_dp_metadata = ( + None if num_tokens_across_dp is None else int(num_tokens_across_dp) + ) + if ( + os.environ.get("VLLM_XPU_CUDAGRAPH_CLEAR_ON_PREFILL", "0") == "1" + and num_reqs > 0 + and prefill_req_count == num_reqs + ): + req_key = "|".join(str(req_id) for req_id in req_ids[:num_reqs]) + if req_key and req_key != self.xpu_cudagraph_last_prefill_clear_req_key: + if hasattr(torch, "xpu"): + torch.xpu.synchronize() + CUDAGraphWrapper.clear_all_graphs() + set_cudagraph_capturing_enabled(True) + self.xpu_decode_cudagraph_replay_count = 0 + self.xpu_cudagraph_last_prefill_clear_req_key = req_key + + self._xpu_timing_step_active = begin_timing_step( + { + "num_reqs": int(num_reqs), + "num_tokens_unpadded": int(num_tokens_unpadded), + "num_tokens_padded": int(num_tokens_padded), + "max_num_scheduled_tokens": int(max_num_scheduled_tokens), + "scheduled_token_counts": scheduled_token_counts, + "scheduled_token_histogram": scheduled_token_histogram, + "scheduled_spec_lengths": scheduled_spec_lengths, + "scheduled_spec_histogram": scheduled_spec_histogram, + "max_scheduled_spec_tokens": int(max_scheduled_spec_tokens), + "computed_token_counts": computed_token_counts, + "decode_req_count": int(decode_req_count), + "prefill_req_count": int(prefill_req_count), + "is_pure_decode": bool(is_pure_decode), + "decode_bucket": decode_bucket, + "cudagraph_mode": str(cudagraph_mode), + "batch_desc_num_tokens": int(batch_desc.num_tokens), + "batch_desc_num_reqs": None + if batch_desc.num_reqs is None + else int(batch_desc.num_reqs), + "num_tokens_across_dp": num_tokens_across_dp_metadata, + "skip_compiled": bool(skip_compiled), + "use_spec_decode": bool(use_spec_decode), + "should_ubatch": bool(should_ubatch), + } + ) + device_type = getattr(self.device, "type", "cuda") + execute_device_timeline: dict[str, object] | None = None + if ( + ( + _XPU_ASYNC_OUTPUT_DEVICE_TIMING + or _XPU_PRE_SAMPLER_BOUNDARY_TIMING + or _XPU_FORWARD_BOUNDARY_SYNC + ) + and self.use_async_scheduling + ): + execute_device_timeline = { + "execute_entry_event": _xpu_async_output_record_timeline_event( + "execute_entry_event", + device_type + ), + "host_execute_entry_ns": time.perf_counter_ns(), + "pre_sampler_device_type": device_type, + "pre_sampler_num_reqs": int(num_reqs), + "pre_sampler_num_tokens_unpadded": int(num_tokens_unpadded), + "pre_sampler_num_tokens_padded": int(num_tokens_padded), + "pre_sampler_is_pure_decode": bool(is_pure_decode), + "pre_sampler_decode_bucket": decode_bucket, + "pre_sampler_cudagraph_mode": str(cudagraph_mode), + "pre_sampler_use_local_argmax": False, + "pre_sampler_broadcast_pp_output": bool(self.broadcast_pp_output), + } # Run the model. # Use persistent buffers for CUDA graphs. @@ -4069,6 +6671,7 @@ class GPUModelRunner( # (wait_for_save + clear metadata) until after draft model runs. defer_kv_connector_finalize = self.speculative_config is not None with ( + timed_region("gpu_model_runner.forward_total"), set_forward_context( attn_metadata, self.vllm_config, @@ -4078,7 +6681,21 @@ class GPUModelRunner( batch_descriptor=batch_desc, ubatch_slices=ubatch_slices_padded, slot_mapping=slot_mappings, - skip_compiled=has_encoder_input, + skip_compiled=skip_compiled, + additional_kwargs={ + "xpu_req_ids": list(req_ids[:num_reqs]), + "xpu_num_prompt_tokens": [ + int(v) + for v in self.input_batch.num_prompt_tokens[:num_reqs] + ], + "xpu_num_computed_tokens": [ + int(v) + for v in self.input_batch.num_computed_tokens_cpu[:num_reqs] + ], + "xpu_cudagraph_mode": str(cudagraph_mode), + "xpu_num_tokens_unpadded": int(num_tokens_unpadded), + "xpu_num_tokens_padded": int(num_tokens_padded), + }, ), record_function_or_nullcontext("gpu_model_runner: forward"), self.maybe_get_kv_connector_output( @@ -4086,6 +6703,17 @@ class GPUModelRunner( defer_finalize=defer_kv_connector_finalize, ) as kv_connector_output, ): + if execute_device_timeline is not None: + execute_device_timeline["forward_start_event"] = ( + _xpu_async_output_record_timeline_event( + "forward_start_event", device_type + ) + ) + execute_device_timeline["host_forward_start_ns"] = ( + time.perf_counter_ns() + ) + if route_overlay_enabled(): + reset_route_overlay_snapshot() model_output = self._model_forward( input_ids=input_ids, positions=positions, @@ -4093,8 +6721,84 @@ class GPUModelRunner( inputs_embeds=inputs_embeds, **model_kwargs, ) + if execute_device_timeline is not None: + execute_device_timeline["forward_end_event"] = ( + _xpu_async_output_record_timeline_event( + "forward_end_event", device_type + ) + ) + execute_device_timeline["host_forward_end_ns"] = ( + time.perf_counter_ns() + ) + route_overlay_snapshot = ( + pop_route_overlay_snapshot() if route_overlay_enabled() else None + ) + should_print_forward, forward_event_id, forward_rank = ( + _xpu_forward_boundary_should_print() + ) + if should_print_forward and execute_device_timeline is not None: + forward_start_event = execute_device_timeline.get( + "forward_start_event" + ) + forward_end_event = execute_device_timeline.get("forward_end_event") + forward_sync_total_start_ns = time.perf_counter_ns() + forward_start_sync_ms = None + forward_end_after_start_sync_ms = None + if forward_start_event is not None: + forward_start_sync_start_ns = time.perf_counter_ns() + forward_start_event.synchronize() + forward_start_sync_ms = _xpu_async_output_ms_since( + forward_start_sync_start_ns + ) + if forward_end_event is not None: + forward_end_sync_start_ns = time.perf_counter_ns() + forward_end_event.synchronize() + forward_end_after_start_sync_ms = _xpu_async_output_ms_since( + forward_end_sync_start_ns + ) + forward_payload = { + "event_id": forward_event_id, + "rank": forward_rank, + "pid": os.getpid(), + "device_type": device_type, + "device": str(self.device), + "local_rank": os.environ.get("LOCAL_RANK"), + "global_rank": os.environ.get("RANK"), + "num_reqs": int(num_reqs), + "num_tokens_unpadded": int(num_tokens_unpadded), + "num_tokens_padded": int(num_tokens_padded), + "max_num_scheduled_tokens": int(max_num_scheduled_tokens), + "is_pure_decode": bool(is_pure_decode), + "decode_bucket": decode_bucket, + "cudagraph_mode": str(cudagraph_mode), + "batch_desc_num_tokens": int(batch_desc.num_tokens), + "batch_desc_num_reqs": None + if batch_desc.num_reqs is None + else int(batch_desc.num_reqs), + "skip_compiled": bool(skip_compiled), + "use_spec_decode": bool(use_spec_decode), + "should_ubatch": bool(should_ubatch), + "forward_start_sync_ms": forward_start_sync_ms, + "forward_end_after_start_sync_ms": ( + forward_end_after_start_sync_ms + ), + "forward_sync_total_ms": _xpu_async_output_ms_since( + forward_sync_total_start_ns + ), + } + if route_overlay_snapshot is not None: + forward_payload["route_overlay"] = route_overlay_snapshot + _xpu_forward_boundary_print(forward_payload) + if ( + os.environ.get("VLLM_XPU_SYNC_AFTER_MODEL_FORWARD", "0") == "1" + and torch.xpu.is_available() + ): + torch.xpu.synchronize() - with record_function_or_nullcontext("gpu_model_runner: postprocess"): + with ( + record_function_or_nullcontext("gpu_model_runner: postprocess"), + timed_region("gpu_model_runner.postprocess_total"), + ): if self.use_aux_hidden_state_outputs: # True when EAGLE 3 is used. hidden_states, aux_hidden_states = model_output @@ -4102,6 +6806,16 @@ class GPUModelRunner( # Common case. hidden_states = model_output aux_hidden_states = None + self._trace_xpu_replay_microscope( + stage="hidden_after_forward", + scheduler_output=scheduler_output, + num_reqs=num_reqs, + batch_desc=batch_desc, + cudagraph_mode=cudagraph_mode, + hidden_states=hidden_states + if torch.is_tensor(hidden_states) + else None, + ) if not self.broadcast_pp_output: # Common case. @@ -4110,24 +6824,218 @@ class GPUModelRunner( assert isinstance(hidden_states, IntermediateTensors) hidden_states.kv_connector_output = kv_connector_output self.kv_connector_output = kv_connector_output + self._finish_xpu_timing_step("intermediate") return hidden_states if self.is_pooling_model: # Return the pooling output. - return self._pool( + output = self._pool( hidden_states, num_scheduled_tokens, num_scheduled_tokens_np, kv_connector_output, ) + self._finish_xpu_timing_step("pool") + return output - sample_hidden_states = hidden_states[logits_indices] - logits = self.model.compute_logits(sample_hidden_states) + trace_tensor( + "gpu_model_runner.hidden_states_before_logits_index", + hidden_states, + ) + trace_tensor("gpu_model_runner.logits_indices", logits_indices) + if execute_device_timeline is not None: + execute_device_timeline["select_sample_hidden_start_event"] = ( + _xpu_async_output_record_timeline_event( + "select_sample_hidden_start_event", device_type + ) + ) + execute_device_timeline["host_select_sample_hidden_start_ns"] = ( + time.perf_counter_ns() + ) + with timed_region("gpu_model_runner.select_sample_hidden"): + sample_hidden_states = _select_sample_hidden_states( + hidden_states, logits_indices + ) + self._trace_xpu_replay_microscope( + stage="sample_hidden", + scheduler_output=scheduler_output, + num_reqs=num_reqs, + batch_desc=batch_desc, + cudagraph_mode=cudagraph_mode, + logits_indices=logits_indices, + sample_hidden_states=sample_hidden_states, + ) + if execute_device_timeline is not None: + execute_device_timeline["select_sample_hidden_end_event"] = ( + _xpu_async_output_record_timeline_event( + "select_sample_hidden_end_event", device_type + ) + ) + execute_device_timeline["host_select_sample_hidden_end_ns"] = ( + time.perf_counter_ns() + ) + if ( + os.environ.get("VLLM_XPU_CLONE_SAMPLE_HIDDEN", "0") == "1" + and sample_hidden_states.device.type == "xpu" + ): + if execute_device_timeline is not None: + execute_device_timeline["clone_sample_hidden_start_event"] = ( + _xpu_async_output_record_timeline_event( + "clone_sample_hidden_start_event", device_type + ) + ) + execute_device_timeline[ + "host_clone_sample_hidden_start_ns" + ] = time.perf_counter_ns() + with timed_region("gpu_model_runner.clone_sample_hidden"): + sample_hidden_states = ( + sample_hidden_states.contiguous().clone() + ) + if execute_device_timeline is not None: + execute_device_timeline["clone_sample_hidden_end_event"] = ( + _xpu_async_output_record_timeline_event( + "clone_sample_hidden_end_event", device_type + ) + ) + execute_device_timeline[ + "host_clone_sample_hidden_end_ns" + ] = time.perf_counter_ns() + trace_tensor( + "gpu_model_runner.sample_hidden_states", sample_hidden_states + ) + if self._can_use_xpu_local_argmax(spec_decode_metadata): + if execute_device_timeline is not None: + execute_device_timeline["pre_sampler_use_local_argmax"] = True + execute_device_timeline["local_argmax_start_event"] = ( + _xpu_async_output_record_timeline_event( + "local_argmax_start_event", device_type + ) + ) + execute_device_timeline["host_local_argmax_start_ns"] = ( + time.perf_counter_ns() + ) + sampled_token_ids = self.model.get_top_tokens(sample_hidden_states) + if execute_device_timeline is not None: + execute_device_timeline["local_argmax_end_event"] = ( + _xpu_async_output_record_timeline_event( + "local_argmax_end_event", device_type + ) + ) + execute_device_timeline["host_local_argmax_end_ns"] = ( + time.perf_counter_ns() + ) + sampled_token_ids_i32 = self._sync_xpu_local_argmax_sampled_token_ids( + sampled_token_ids.to(torch.int32) + ) + self._xpu_local_argmax_sampled_token_ids = ( + sampled_token_ids_i32.unsqueeze(-1) + ) + trace_tensor( + "gpu_model_runner.local_argmax_sampled_token_ids", + self._xpu_local_argmax_sampled_token_ids, + ) + self._trace_xpu_replay_microscope( + stage="local_argmax_sampled", + scheduler_output=scheduler_output, + num_reqs=num_reqs, + batch_desc=batch_desc, + cudagraph_mode=cudagraph_mode, + sample_hidden_states=sample_hidden_states, + extra={ + "sampled_token_ids": sampled_token_ids_i32.detach() + .to("cpu") + .tolist() + }, + ) + logits = None + else: + self._xpu_local_argmax_sampled_token_ids = None + if execute_device_timeline is not None: + execute_device_timeline["compute_logits_start_event"] = ( + _xpu_async_output_record_timeline_event( + "compute_logits_start_event", device_type + ) + ) + execute_device_timeline["host_compute_logits_start_ns"] = ( + time.perf_counter_ns() + ) + with timed_region("gpu_model_runner.compute_logits"): + logits = self.model.compute_logits(sample_hidden_states) + if execute_device_timeline is not None: + execute_device_timeline["compute_logits_end_event"] = ( + _xpu_async_output_record_timeline_event( + "compute_logits_end_event", device_type + ) + ) + execute_device_timeline["host_compute_logits_end_ns"] = ( + time.perf_counter_ns() + ) + self._trace_xpu_replay_microscope( + stage="logits_after_compute", + scheduler_output=scheduler_output, + num_reqs=num_reqs, + batch_desc=batch_desc, + cudagraph_mode=cudagraph_mode, + logits=logits, + ) + trace_tensor("gpu_model_runner.logits_after_compute", logits) else: # Rare case. assert not self.is_pooling_model - sample_hidden_states = hidden_states[logits_indices] + trace_tensor( + "gpu_model_runner.hidden_states_before_logits_index", + hidden_states, + ) + trace_tensor("gpu_model_runner.logits_indices", logits_indices) + if execute_device_timeline is not None: + execute_device_timeline["select_sample_hidden_start_event"] = ( + _xpu_async_output_record_timeline_event( + "select_sample_hidden_start_event", device_type + ) + ) + execute_device_timeline["host_select_sample_hidden_start_ns"] = ( + time.perf_counter_ns() + ) + with timed_region("gpu_model_runner.select_sample_hidden"): + sample_hidden_states = _select_sample_hidden_states( + hidden_states, logits_indices + ) + if execute_device_timeline is not None: + execute_device_timeline["select_sample_hidden_end_event"] = ( + _xpu_async_output_record_timeline_event( + "select_sample_hidden_end_event", device_type + ) + ) + execute_device_timeline["host_select_sample_hidden_end_ns"] = ( + time.perf_counter_ns() + ) + if ( + os.environ.get("VLLM_XPU_CLONE_SAMPLE_HIDDEN", "0") == "1" + and sample_hidden_states.device.type == "xpu" + ): + if execute_device_timeline is not None: + execute_device_timeline["clone_sample_hidden_start_event"] = ( + _xpu_async_output_record_timeline_event( + "clone_sample_hidden_start_event", device_type + ) + ) + execute_device_timeline[ + "host_clone_sample_hidden_start_ns" + ] = time.perf_counter_ns() + with timed_region("gpu_model_runner.clone_sample_hidden"): + sample_hidden_states = ( + sample_hidden_states.contiguous().clone() + ) + if execute_device_timeline is not None: + execute_device_timeline["clone_sample_hidden_end_event"] = ( + _xpu_async_output_record_timeline_event( + "clone_sample_hidden_end_event", device_type + ) + ) + execute_device_timeline[ + "host_clone_sample_hidden_end_ns" + ] = time.perf_counter_ns() if not get_pp_group().is_last_rank: all_gather_tensors = { "residual": not is_residual_scattered_for_sp( @@ -4141,7 +7049,27 @@ class GPUModelRunner( ) logits = None else: - logits = self.model.compute_logits(sample_hidden_states) + if execute_device_timeline is not None: + execute_device_timeline["compute_logits_start_event"] = ( + _xpu_async_output_record_timeline_event( + "compute_logits_start_event", device_type + ) + ) + execute_device_timeline["host_compute_logits_start_ns"] = ( + time.perf_counter_ns() + ) + with timed_region("gpu_model_runner.compute_logits"): + logits = self.model.compute_logits(sample_hidden_states) + if execute_device_timeline is not None: + execute_device_timeline["compute_logits_end_event"] = ( + _xpu_async_output_record_timeline_event( + "compute_logits_end_event", device_type + ) + ) + execute_device_timeline["host_compute_logits_end_ns"] = ( + time.perf_counter_ns() + ) + trace_tensor("gpu_model_runner.logits_after_compute", logits) model_output_broadcast_data: dict[str, Any] = {} if logits is not None: @@ -4164,7 +7092,17 @@ class GPUModelRunner( ec_connector_output, cudagraph_stats, slot_mappings, + execute_device_timeline, ) + if execute_device_timeline is not None: + execute_device_timeline["execute_state_ready_event"] = ( + _xpu_async_output_record_timeline_event( + "execute_state_ready_event", device_type + ) + ) + execute_device_timeline["host_execute_state_ready_ns"] = ( + time.perf_counter_ns() + ) self.kv_connector_output = kv_connector_output # Now the batch has been launched we can wait for corrections from the @@ -4208,6 +7146,7 @@ class GPUModelRunner( ec_connector_output, cudagraph_stats, slot_mappings, + device_timeline, ) = self.execute_model_state # Clear ephemeral state. self.execute_model_state = None @@ -4218,12 +7157,72 @@ class GPUModelRunner( scheduler_output, grammar_output, self.input_batch, logits ) - with record_function_or_nullcontext("gpu_model_runner: sample"): + device_type = getattr(self.device, "type", "cuda") + if ( + (_XPU_ASYNC_OUTPUT_DEVICE_TIMING or _XPU_PRE_SAMPLER_BOUNDARY_TIMING) + and self.use_async_scheduling + ): + if device_timeline is None: + device_timeline = {} + device_timeline["sample_start_event"] = ( + _xpu_async_output_record_timeline_event( + "sample_start_event", device_type + ) + ) + device_timeline["host_sample_start_ns"] = time.perf_counter_ns() + else: + device_timeline = None + + self._trace_xpu_replay_microscope( + stage="pre_sample", + scheduler_output=scheduler_output, + num_reqs=self.input_batch.num_reqs, + logits=logits, + ) + with ( + record_function_or_nullcontext("gpu_model_runner: sample"), + timed_region("gpu_model_runner.sample_total"), + ): sampler_output = self._sample(logits, spec_decode_metadata) + self._trace_xpu_replay_microscope( + stage="sampler_output", + scheduler_output=scheduler_output, + num_reqs=self.input_batch.num_reqs, + logits=logits, + sampler_output=sampler_output, + ) + if device_timeline is not None: + sampler_timeline = getattr(sampler_output, "_xpu_sampler_timeline", None) + if isinstance(sampler_timeline, dict): + device_timeline.update(sampler_timeline) + device_timeline["sample_end_event"] = ( + _xpu_async_output_record_timeline_event( + "sample_end_event", device_type + ) + ) + device_timeline["host_sample_end_ns"] = time.perf_counter_ns() + + if ( + self.use_async_scheduling + and os.environ.get("VLLM_XPU_ASYNC_CLONE_SAMPLED_TOKEN_IDS", "0") == "1" + and sampler_output.sampled_token_ids.device.type == "xpu" + ): + # XPU async scheduling reuses the sampled-token tensor both as the + # next decode input and for an overlapped D2H output copy. Keep a + # private device allocation to avoid sampler-buffer lifetime/order + # issues while staying on the GPU fast path. + sampler_output.sampled_token_ids = sampler_output.sampled_token_ids.clone() self._update_states_after_model_execute( sampler_output.sampled_token_ids, scheduler_output ) + if device_timeline is not None: + device_timeline["state_update_end_event"] = ( + _xpu_async_output_record_timeline_event( + "state_update_end_event", device_type + ) + ) + device_timeline["host_state_update_end_ns"] = time.perf_counter_ns() if self.use_async_scheduling: pp = get_pp_group() # For torchrun external_launcher PP mode with broadcast_pp_output=True, @@ -4253,6 +7252,36 @@ class GPUModelRunner( spec_decode_common_attn_metadata, slot_mappings, ) + if os.environ.get("VLLM_XPU_TRACE_NGRAM_PP"): + if torch.is_tensor(self._draft_token_ids): + draft_summary = { + "shape": tuple(self._draft_token_ids.shape), + "dtype": str(self._draft_token_ids.dtype), + } + else: + draft_summary = [ + len(toks) for toks in self._draft_token_ids[:4] + ] + if torch.is_tensor(sampled_token_ids): + sampled_summary = { + "shape": tuple(sampled_token_ids.shape), + "dtype": str(sampled_token_ids.dtype), + } + else: + sampled_summary = [ + len(toks) for toks in sampled_token_ids[:4] + ] + logger.warning( + "XPU ngram proposed pp_rank=%s first=%s last=%s " + "method=%s sampled=%s draft=%s req_ids=%s", + get_pp_group().rank_in_group, + get_pp_group().is_first_rank, + get_pp_group().is_last_rank, + spec_config.method, + sampled_summary, + draft_summary, + self.input_batch.req_ids, + ) self._copy_draft_token_ids_to_cpu(scheduler_output) spec_config = self.speculative_config @@ -4331,21 +7360,29 @@ class GPUModelRunner( self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True) with record_function_or_nullcontext("gpu_model_runner: bookkeep"): - ( - num_nans_in_logits, - logprobs_lists, - valid_sampled_token_ids, - prompt_logprobs_dict, - req_ids_output_copy, - req_id_to_index_output_copy, - invalid_req_indices, - ) = self._bookkeeping_sync( - scheduler_output, - sampler_output, - logits, - hidden_states, - scheduler_output.total_num_scheduled_tokens, + with timed_region("gpu_model_runner.bookkeeping_sync"): + ( + num_nans_in_logits, + logprobs_lists, + valid_sampled_token_ids, + prompt_logprobs_dict, + req_ids_output_copy, + req_id_to_index_output_copy, + invalid_req_indices, + ) = self._bookkeeping_sync( + scheduler_output, + sampler_output, + logits, + hidden_states, + scheduler_output.total_num_scheduled_tokens, + ) + if device_timeline is not None: + device_timeline["bookkeeping_end_event"] = ( + _xpu_async_output_record_timeline_event( + "bookkeeping_end_event", device_type + ) ) + device_timeline["host_bookkeeping_end_ns"] = time.perf_counter_ns() if propose_drafts_after_bookkeeping: # ngram and other speculative decoding methods use the sampled @@ -4386,12 +7423,23 @@ class GPUModelRunner( num_nans_in_logits=num_nans_in_logits, cudagraph_stats=cudagraph_stats, ) + if device_timeline is not None: + device_timeline["pre_async_wrap_event"] = ( + _xpu_async_output_record_timeline_event( + "pre_async_wrap_event", device_type + ) + ) + device_timeline["host_pre_async_wrap_ns"] = time.perf_counter_ns() if not self.use_async_scheduling: + self._finish_xpu_timing_step("ok") return output - with record_function_or_nullcontext( - "gpu_model_runner: AsyncGPUModelRunnerOutput" + with ( + record_function_or_nullcontext( + "gpu_model_runner: AsyncGPUModelRunnerOutput" + ), + timed_region("gpu_model_runner.async_output_wrap"), ): async_output = AsyncGPUModelRunnerOutput( model_runner_output=output, @@ -4400,6 +7448,12 @@ class GPUModelRunner( invalid_req_indices=invalid_req_indices, async_output_copy_stream=self._get_or_create_async_output_copy_stream(), vocab_size=self.input_batch.vocab_size, + sampled_token_ids_cpu_buffer=( + self._get_async_sampled_token_ids_cpu_buffer( + sampler_output.sampled_token_ids + ) + ), + device_timeline=device_timeline, ) with record_function_or_nullcontext( "gpu_model_runner: set_async_sampled_token_ids" @@ -4411,6 +7465,7 @@ class GPUModelRunner( async_output.async_copy_ready_event, ) + self._finish_xpu_timing_step("ok") return async_output def _pp_broadcast_prev_sampled_token_ids( @@ -4459,8 +7514,28 @@ class GPUModelRunner( def take_draft_token_ids(self) -> DraftTokenIds | None: if not self.num_spec_tokens or not self._draft_token_req_ids: + if os.environ.get("VLLM_XPU_TRACE_NGRAM_PP"): + logger.warning( + "XPU ngram take_draft none pp_rank=%s first=%s last=%s " + "num_spec_tokens=%s has_req_ids=%s", + get_pp_group().rank_in_group, + get_pp_group().is_first_rank, + get_pp_group().is_last_rank, + self.num_spec_tokens, + bool(self._draft_token_req_ids), + ) return None draft_token_ids, req_ids = self._get_draft_token_ids_cpu() + if os.environ.get("VLLM_XPU_TRACE_NGRAM_PP"): + logger.warning( + "XPU ngram take_draft pp_rank=%s first=%s last=%s " + "req_ids=%s draft_lens=%s", + get_pp_group().rank_in_group, + get_pp_group().is_first_rank, + get_pp_group().is_last_rank, + req_ids, + [len(toks) for toks in draft_token_ids[:4]], + ) return DraftTokenIds(req_ids, draft_token_ids) def _copy_draft_token_ids_to_cpu( @@ -5570,6 +8645,11 @@ class GPUModelRunner( batch_descriptor=batch_desc, ubatch_slices=ubatch_slices_padded, slot_mapping=slot_mappings, + # Keep profiling/capture on the normal compiled path. The + # XPU prefill bypass is a live-runtime repair knob; applying + # it here can feed synthetic max_tokens + 1 profile shapes + # into the piecewise backend before those ranges exist. + skip_compiled=False, ), ): outputs = self.model( @@ -6144,6 +9224,18 @@ class GPUModelRunner( start_time = time.perf_counter() + if os.environ.get("VLLM_MINIMAX_QK_NORM_PRECAPTURE_SANITIZE", "0") == "1": + raw_model = self.get_model() + sanitize_qk_norms = getattr( + raw_model, "sanitize_qk_norm_weights_for_capture", None + ) + if callable(sanitize_qk_norms): + stats = sanitize_qk_norms() + torch.accelerator.synchronize() + logger.info( + "MiniMax q/k RMSNorm pre-capture sanitize: %s", stats + ) + # Trigger CUDA graph capture for specific shapes. # Capture the large shapes first so that the smaller shapes # can reuse the memory pool allocated for the large shapes. @@ -7041,12 +10133,41 @@ class GPUModelRunner( # this is in the critical path of every single model # forward loop, this has caused perf issue for a disagg # setup. + if ( + sampled_token_ids.device.type == "xpu" + and os.environ.get("VLLM_XPU_SYNC_TO_LIST", "0") == "1" + ): + if os.environ.get("VLLM_XPU_SYNC_TO_LIST_BEFORE_COPY", "1") != "0": + torch.xpu.synchronize() + return sampled_token_ids.to("cpu").tolist() pinned = self.sampled_token_ids_pinned_cpu[: sampled_token_ids.shape[0]] pinned.copy_(sampled_token_ids, non_blocking=True) self.transfer_event.record() self.transfer_event.synchronize() return pinned.tolist() + def _get_async_sampled_token_ids_cpu_buffer( + self, sampled_token_ids: torch.Tensor + ) -> torch.Tensor | None: + if os.environ.get("VLLM_XPU_REUSE_ASYNC_OUTPUT_COPY_BUFFER", "0") != "1": + return None + buffers = self.async_sampled_token_ids_pinned_cpu_buffers + if not buffers or sampled_token_ids.device.type != "xpu": + return None + next_index = self.async_sampled_token_ids_pinned_cpu_buffer_index + self.async_sampled_token_ids_pinned_cpu_buffer_index = ( + next_index + 1 + ) % len(buffers) + buffer = buffers[next_index] + if buffer.dtype != sampled_token_ids.dtype: + return None + if ( + buffer.shape[0] < sampled_token_ids.shape[0] + or buffer.shape[1] < sampled_token_ids.shape[1] + ): + return None + return buffer + def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]: """ Get encoder timing stats for all requests and clear the registry.