diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 7b9478035ece..4bdeb7e45a20 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -1185,11 +1185,35 @@ def set_splitting_ops_for_v1( ) self.cudagraph_mode = CUDAGraphMode.FULL + # With worst-token dispatch (UCCL ht-cudagraph-worst-tokens kernels), + # DeepEP HT dispatch+combine are stream-capturable: no host count + # sync, static worst-case recv shapes, device-side per-expert counts. + # Capture decode steps fully; prefill stays eager. + if ( + all2all_backend == "deepep_high_throughput" + and data_parallel_size > 1 + and self.cudagraph_mode != CUDAGraphMode.NONE + and envs.VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH + ): + if self.cudagraph_mode.has_full_cudagraphs(): + logger.info( + "DeepEP high-throughput worst-token dispatch enabled; " + "capturing full decode CUDA graphs." + ) + self.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY + else: + logger.warning_once( + "VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH requires full decode " + "CUDA graphs; setting cudagraph_mode to NONE." + ) + self.cudagraph_mode = CUDAGraphMode.NONE + # Disable CUDA graphs for DeepEP high-throughput since its not CG compatible if ( all2all_backend == "deepep_high_throughput" and data_parallel_size > 1 and self.cudagraph_mode != CUDAGraphMode.NONE + and not envs.VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH ): # TODO: Piecewise Cuda graph might be enabled # if torch compile cache key issue fixed diff --git a/vllm/envs.py b/vllm/envs.py index b2c5f22567fa..1114c180f533 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -151,6 +151,7 @@ VLLM_USE_STANDALONE_COMPILE: bool = True VLLM_ENABLE_PREGRAD_PASSES: bool = True VLLM_USE_BREAKABLE_CUDAGRAPH: bool = False + VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH: bool = False VLLM_DP_MASTER_IP: str = "" VLLM_DP_MASTER_PORT: int = 0 VLLM_RANDOMIZE_DP_DUMMY_INPUTS: bool = False @@ -719,6 +720,12 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_USE_BREAKABLE_CUDAGRAPH": lambda: ( os.environ.get("VLLM_USE_BREAKABLE_CUDAGRAPH", "0") == "1" ), + # Experimental: DeepEP high-throughput dispatch with worst-case-sized + # recv buffers and no host count sync (CUDA-graph capturable; requires + # the UCCL ht-cudagraph-worst-tokens kernels). + "VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH": lambda: ( + os.environ.get("VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH", "0") == "1" + ), # Debug pattern matching inside custom passes. # Should be set to the fx.Node name (e.g. 'getitem_34' or 'scaled_mm_3'). "VLLM_PATTERN_MATCH_DEBUG": lambda: os.environ.get( diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py index 9ca94521363e..b81e8403bf8e 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py @@ -5,7 +5,13 @@ import deep_ep import torch +import vllm.envs as envs import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.config import CUDAGraphMode +from vllm.forward_context import ( + get_forward_context, + is_forward_context_available, +) from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceContiguous, @@ -65,9 +71,68 @@ def __init__( # micro-batch to avoid races between threads. self.handles = [None, None] + # Lamport-clock counters for host-issued dispatch/combine calls, + # reported in the timeout annotations below. Healthy EP ranks + # advance these in lockstep (graph REPLAYS bump no rank's counter + # since no Python runs, and the runtime mode is synced across DP + # ranks per step), so on a collective timeout, comparing values + # across ranks' logs separates "all ranks in the same stuck + # collective" (transport) from "ranks in different collectives" + # (scheduling desync) — and seq mod num-moe-layers localizes the + # layer. + self.dispatch_seq = 0 + self.combine_seq = 0 + # From https://github.com/deepseek-ai/DeepEP/blob/9fe9021f29c9083cd1808ab36b740208524d9f63/deep_ep/buffer.py#L164 self.available_rank_configs = [2, 4, 8, 16, 24, 32, 64, 128, 144, 160] + def _step_debug(self) -> str: + try: + if is_forward_context_available(): + ctx = get_forward_context() + return ( + f"mode={ctx.cudagraph_runtime_mode} " + f"descriptor={ctx.batch_descriptor}" + ) + except Exception: + pass + return "no-forward-context" + + def _num_worst_tokens(self, tokens: torch.Tensor) -> int: + """Worst-case recv size for CUDA-graph-safe dispatch. + + When non-zero, dispatch skips its host count sync, returns + worst-case-padded recv tensors with per-expert counts in a device + tensor, and the dispatch+combine round trip becomes + stream-capturable (requires the UCCL ht-cudagraph-worst-tokens + kernels). + + Sized from the dispatched tensor itself: DP padding makes its row + count uniform across ranks, so it equals the step's padded token + count for whole-batch steps and this microbatch's padded slice + under DBO ubatching (the forward context's batch descriptor covers + the whole step there, which would oversize ubatch recv buffers 2x). + Eager steps (prefill, mixed, uncaptured shapes) return 0 and keep + the host-synced path — no padding overhead where there is no graph + to serve.""" + if not envs.VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH: + return 0 + if not is_forward_context_available(): + return 0 + ctx = get_forward_context() + # UBatchWrapper passes cudagraph_runtime_mode=NONE into the + # per-ubatch forward contexts while CAPTURING a full ubatched + # graph, so the context mode alone misses exactly the case that + # must not record the host-synced protocol (its CPU poll waits on + # a kernel that never executes under capture). Check the stream's + # capture status directly as well. + if ( + ctx.cudagraph_runtime_mode != CUDAGraphMode.FULL + and not torch.cuda.is_current_stream_capturing() + ): + return 0 + return tokens.size(0) * self.num_dispatchers_ + def num_dispatchers(self) -> int: return self.num_dispatchers_ @@ -136,15 +201,18 @@ def _do_dispatch( if has_scales: token_data = (tokens, token_scales) - ( - token_data, - expert_topk_ids, - expert_topk_weights, - expert_num_tokens_per_expert_list, - handle, - event, - ) = self.buffer.dispatch( - x=token_data, + num_worst_tokens = self._num_worst_tokens(tokens) + self.dispatch_seq += 1 + try: + ( + token_data, + expert_topk_ids, + expert_topk_weights, + expert_num_tokens_per_expert_list, + handle, + event, + ) = self.buffer.dispatch( + x=token_data, handle=None, num_tokens_per_rank=num_tokens_per_rank, num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, @@ -154,12 +222,29 @@ def _do_dispatch( topk_weights=rank_topk_weights, # expert_alignment rounds the number of tokens per expert # to this value. - expert_alignment=1, - config=self._get_dispatch_config(), - previous_event=previous_event, - async_finish=self.async_prepare and not dbo_enabled(), - allocate_on_comm_stream=False, - ) + expert_alignment=1, + num_worst_tokens=num_worst_tokens, + config=self._get_dispatch_config(), + previous_event=previous_event, + # With worst-token graphs in play, ALL dispatches must join + # back to the compute stream: a graph replay launched on the + # compute stream does not wait on the live comm stream, so an + # async eager dispatch still in flight would race the captured + # dispatch kernels on the shared ring buffers. + async_finish=( + self.async_prepare + and not dbo_enabled() + and not envs.VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH + ), + allocate_on_comm_stream=False, + ) + except RuntimeError as exc: + raise RuntimeError( + f"{exc} [deepep-ht dispatch debug: " + f"seq={self.dispatch_seq} ubatch={dbo_current_ubatch_id()} " + f"tokens={tokens.shape[0]} worst={num_worst_tokens} " + f"{self._step_debug()}]" + ) from exc # record the handle for this ubatch a2a_idx = dbo_current_ubatch_id() @@ -187,7 +272,7 @@ def _receiver( token_data: tuple[torch.Tensor, torch.Tensor] | torch.Tensor, expert_topk_ids: torch.Tensor | None, num_experts: int, - expert_num_tokens_per_expert_list: list[int], + expert_num_tokens_per_expert_list: list[int] | torch.Tensor, expert_topk_weights: torch.Tensor | None, a1_scale: torch.Tensor | None, quant_config: FusedMoEQuantConfig, @@ -219,12 +304,21 @@ def _receiver( expert_topk_ids + self.rank_expert_offset, ) - # Makes a GPU-CPU copy. - # TODO (varun): Maybe it is better to re-compute the expert_num_tokens - # on GPU. - expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list( - expert_num_tokens_per_expert_list, device=expert_x.device - ) + if isinstance(expert_num_tokens_per_expert_list, torch.Tensor): + # Worst-token (CUDA-graph) dispatch: counts arrive as a device + # tensor written by the notify kernel; no host copy exists and + # none may be made on this path. + expert_tokens_meta = mk.ExpertTokensMetadata( + expert_num_tokens=expert_num_tokens_per_expert_list, + expert_num_tokens_cpu=None, + ) + else: + # Makes a GPU-CPU copy. + # TODO (varun): Maybe it is better to re-compute the + # expert_num_tokens on GPU. + expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list( + expert_num_tokens_per_expert_list, device=expert_x.device + ) # * For non-block quant, dispatch in b16 and quantize now as # DeepEP kernels only support dispatching block scales. @@ -364,16 +458,28 @@ def _finalize( assert fused_expert_output.dtype == torch.bfloat16, ( f"Expected fused_expert_output bfloat16, got {fused_expert_output.dtype}" ) - combined_x, _, event = self.buffer.combine( + self.combine_seq += 1 + try: + combined_x, _, event = self.buffer.combine( # HT combine only supports BF16 x=fused_expert_output, handle=handle, topk_weights=None, config=self._get_combine_config(), previous_event=previous_event, - async_finish=do_async and not dbo_enabled(), + # Same ring-race hazard as dispatch: combine's comm-stream work + # must be joined before a subsequent graph replay can launch. + async_finish=do_async + and not dbo_enabled() + and not envs.VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH, allocate_on_comm_stream=False, - ) + ) + except RuntimeError as exc: + raise RuntimeError( + f"{exc} [deepep-ht combine debug: " + f"seq={self.combine_seq} ubatch={dbo_current_ubatch_id()} " + f"tokens={fused_expert_output.shape[0]} {self._step_debug()}]" + ) from exc dbo_switch_to_compute() diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index f36dc8f17629..a2ebee35d7fc 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -106,10 +106,13 @@ def build( ) -> CompressorMetadata: query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu num_reqs = common_attn_metadata.num_reqs + num_tokens = common_attn_metadata.num_actual_tokens query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory() - token_to_req_indices = self.token_to_req_indices[: x.shape[0]] - token_to_req_indices.copy_(x, non_blocking=True) + assert x.shape[0] <= num_tokens + self.token_to_req_indices[:num_tokens].fill_(0) + self.token_to_req_indices[: x.shape[0]].copy_(x, non_blocking=True) + token_to_req_indices = self.token_to_req_indices[:num_tokens] return CompressorMetadata( block_table=common_attn_metadata.block_table_tensor.clamp_(min=0), slot_mapping=common_attn_metadata.slot_mapping, diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2870ec9a15c0..a006fa9930a5 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -14,7 +14,6 @@ has_deep_gemm, ) from vllm.utils.math_utils import cdiv -from vllm.utils.platform_utils import num_compute_units from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -28,6 +27,7 @@ ) from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec from vllm.v1.worker.cp_utils import get_total_cp_world_size +from vllm.v1.worker.sm_control import get_ubatch_compute_sms logger = init_logger(__name__) @@ -107,6 +107,10 @@ def split_indexer_prefill_chunks( end += 1 req_slice = slice(start + request_offset, end + request_offset) + # Zero-query rows can appear from padded requests in full-CG/ubatch metadata. + # They produce no logits/indexer work, so skip the no-op chunk. + if chunk_m == 0: + continue max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else chunk_m for q_off in range(0, chunk_m, max_q): sub_m = min(max_q, chunk_m - q_off) @@ -277,8 +281,7 @@ def __init__(self, *args, **kwargs): or not current_platform.is_device_capability_family(100) ) and next_n not in self.natively_supported_next_n_fp4 - sm_count = num_compute_units(self.device.index) - self.num_sms = sm_count + self.num_sms = get_ubatch_compute_sms(self.vllm_config, self.device.index) self.offsets_buffer = torch.arange( next_n, device=self.device, dtype=torch.int32 @@ -477,12 +480,17 @@ def build( seq_lens = common_attn_metadata.seq_lens slot_mapping = common_attn_metadata.slot_mapping block_table = common_attn_metadata.block_table_tensor + has_prefilling_rows = ( + common_attn_metadata.is_prefilling is not None + and torch.any(common_attn_metadata.is_prefilling).item() + ) num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( split_decodes_and_prefills( common_attn_metadata, decode_threshold=self.reorder_batch_threshold, require_uniform=not self.use_flattening, + treat_short_extends_as_decodes=not has_prefilling_rows, ) ) diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index f0e444e493c4..c9b1b82c755d 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -288,10 +288,13 @@ def build( # NOTE: Ensure all metadata tensors maintain fixed memory addresses # for CUDA graph compatibility. + num_tokens = common_attn_metadata.num_actual_tokens query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory() - token_to_req_indices = self.token_to_req_indices[: x.shape[0]] - token_to_req_indices.copy_(x, non_blocking=True) + assert x.shape[0] <= num_tokens + self.token_to_req_indices[:num_tokens].fill_(0) + self.token_to_req_indices[: x.shape[0]].copy_(x, non_blocking=True) + token_to_req_indices = self.token_to_req_indices[:num_tokens] is_valid_token = self.is_valid_token[: slot_mapping.shape[0]] is_valid_token.copy_(slot_mapping >= 0) diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index b73d17e8e5cc..9b94064ffe27 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -480,9 +480,9 @@ def split_decodes_prefills_and_extends( num_extend_tokens: The number of tokens in the extend requests. num_prefill_tokens: The number of tokens in the prefill requests. """ - max_query_len = common_attn_metadata.max_query_len - num_reqs = common_attn_metadata.num_reqs - num_tokens = common_attn_metadata.num_actual_tokens + max_query_len = int(common_attn_metadata.max_query_len) + num_reqs = int(common_attn_metadata.num_reqs) + num_tokens = int(common_attn_metadata.num_actual_tokens) query_start_loc = common_attn_metadata.query_start_loc_cpu # Upper bound is exact for prefill rows; decode rows still satisfy # seq_len > query_len under the optimistic bound, so `seq_lens == @@ -496,10 +496,10 @@ def split_decodes_prefills_and_extends( query_lens = query_start_loc[1:] - query_start_loc[:-1] is_prefill_or_extend = query_lens > decode_threshold is_prefill = (seq_lens == query_lens) & is_prefill_or_extend - first_extend = is_prefill_or_extend.int().argmax(dim=-1).item() - first_prefill = is_prefill.int().argmax(dim=-1).item() + first_extend = int(is_prefill_or_extend.int().argmax(dim=-1).item()) + first_prefill = int(is_prefill.int().argmax(dim=-1).item()) num_decodes = first_extend - num_decode_tokens = query_start_loc[first_extend].item() + num_decode_tokens = int(query_start_loc[first_extend].item()) if not torch.any(is_prefill_or_extend): return (num_decodes, 0, 0, num_decode_tokens, 0, 0) @@ -518,8 +518,8 @@ def split_decodes_prefills_and_extends( num_extends = first_prefill - num_decodes num_prefills = num_reqs - first_prefill - num_prefill_tokens = num_tokens - query_start_loc[first_prefill] - num_extend_tokens = num_prefill_or_extend_tokens - num_prefill_tokens + num_prefill_tokens = int(num_tokens - query_start_loc[first_prefill].item()) + num_extend_tokens = int(num_prefill_or_extend_tokens - num_prefill_tokens) return ( num_decodes, num_extends, @@ -560,9 +560,9 @@ def split_decodes_and_prefills( num_decode_tokens: The number of tokens in the decode requests. num_prefill_tokens: The number of tokens in the prefill requests. """ - max_query_len = common_attn_metadata.max_query_len - num_reqs = common_attn_metadata.num_reqs - num_tokens = common_attn_metadata.num_actual_tokens + max_query_len = int(common_attn_metadata.max_query_len) + num_reqs = int(common_attn_metadata.num_reqs) + num_tokens = int(common_attn_metadata.num_actual_tokens) query_start_loc = common_attn_metadata.query_start_loc_cpu if ( @@ -594,11 +594,11 @@ def split_decodes_and_prefills( if not torch.any(is_prefill): return num_reqs, 0, num_tokens, 0 - first_prefill = is_prefill.int().argmax(dim=-1).item() + first_prefill = int(is_prefill.int().argmax(dim=-1).item()) num_decodes = first_prefill num_prefills = num_reqs - num_decodes - num_decode_tokens = query_start_loc[first_prefill].item() - num_prefill_tokens = num_tokens - num_decode_tokens + num_decode_tokens = int(query_start_loc[first_prefill].item()) + num_prefill_tokens = int(num_tokens - num_decode_tokens) return (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens) diff --git a/vllm/v1/worker/dp_utils.py b/vllm/v1/worker/dp_utils.py index e7c6d81a9929..ed8cded416a6 100644 --- a/vllm/v1/worker/dp_utils.py +++ b/vllm/v1/worker/dp_utils.py @@ -9,7 +9,7 @@ from vllm.logger import init_logger from vllm.v1.worker.ubatch_utils import ( check_ubatch_thresholds, - is_last_ubatch_empty, + has_empty_ubatch, ) logger = init_logger(__name__) @@ -62,11 +62,11 @@ def _post_process_ubatch(tensor: torch.Tensor, num_ubatches: int) -> bool: should_ubatch: bool = bool(torch.all(tensor[2] == 1).item()) if not should_ubatch: return False - # If the DP ranks are planning to ubatch, make sure that - # there are no "empty" second ubatches + # If the DP ranks are planning to ubatch, make sure every microbatch + # has at least one real token. orig_min_num_tokens = int(orig_num_tokens_tensor.min().item()) padded_max_num_tokens = int(padded_num_tokens_tensor.max().item()) - if is_last_ubatch_empty(orig_min_num_tokens, padded_max_num_tokens, num_ubatches): + if has_empty_ubatch(orig_min_num_tokens, padded_max_num_tokens, num_ubatches): logger.debug( "Aborting ubatching %s %s", orig_min_num_tokens, padded_max_num_tokens ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index d51bf2284096..6cc7ee0e44b8 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -2435,7 +2435,12 @@ def _build_attn_group_metadata( for attn_gid in range(len(self.attn_groups[kv_cache_gid])): if ubatch_slices is not None: - for ubid, _cm in enumerate(split_attn_metadata(ubatch_slices, cm)): + split_cms = split_attn_metadata( + ubatch_slices, + cm, + allow_clone=not for_cudagraph_capture, + ) + for ubid, _cm in enumerate(split_cms): _build_attn_group_metadata(kv_cache_gid, attn_gid, _cm, ubid) else: @@ -5731,7 +5736,11 @@ def _dummy_run( cum_num_tokens = self._get_cumsum_and_arange( num_scheduled_tokens, self.query_pos.np ) + self.query_start_loc.np[0] = 0 self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens + self.query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1].fill( + cum_num_tokens[-1] + ) self.query_start_loc.copy_to_gpu() # Sync block table CPU->GPU so cleared rows from @@ -5744,7 +5753,8 @@ def _dummy_run( attn_metadata, _ = self._build_attention_metadata( num_tokens=num_tokens_unpadded, num_tokens_padded=num_tokens_padded if pad_attn else None, - num_reqs=num_reqs_padded, + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded if pad_attn else None, max_query_len=max_query_len, ubatch_slices=(ubatch_slices_padded if pad_attn else ubatch_slices), for_cudagraph_capture=is_graph_capturing, diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py index 657fc8267345..94e37f311430 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/vllm/v1/worker/gpu_ubatch_wrapper.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import threading +import time from collections.abc import Callable from dataclasses import dataclass from typing import Any @@ -11,7 +12,6 @@ import vllm.envs as envs from vllm.compilation.cuda_graph import CUDAGraphWrapper from vllm.config import CUDAGraphMode, VllmConfig -from vllm.distributed import get_ep_group from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id from vllm.forward_context import ( DPMetadata, @@ -26,7 +26,12 @@ from vllm.utils.deep_gemm import set_num_sms as deep_gemm_set_num_sms from vllm.utils.import_utils import has_deep_gemm from vllm.utils.platform_utils import num_compute_units +from vllm.v1.worker.sm_control import ( + get_all2all_manager_for_sm_control, + get_ubatch_comm_sms, +) from vllm.v1.worker.ubatching import UBatchContext, make_ubatch_contexts +from vllm.v1.worker.ubatch_utils import ensure_tensor_alignment logger = init_logger(__name__) @@ -139,6 +144,37 @@ def __init__( self.device = device self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG" self._runnable_str = str(runnable) if self.is_debugging_mode else None + self._mode_counts = { + "ubatched_replay": 0, + "ubatched_eager": 0, + "ubatched_capture": 0, + "plain_eager": 0, + "plain_graph": 0, + } + self._mode_last_log = time.monotonic() + + _MODE_LOG_INTERVAL_S = 30.0 + + def _count_mode(self, mode: str) -> None: + self._mode_counts[mode] += 1 + now = time.monotonic() + if now - self._mode_last_log < self._MODE_LOG_INTERVAL_S: + return + + self._mode_last_log = now + counts = self._mode_counts + logger.info( + "[tbo-modes] last %ds: ubatched_replay=%d ubatched_eager=%d " + "ubatched_capture=%d plain_eager=%d plain_graph=%d", + int(self._MODE_LOG_INTERVAL_S), + counts["ubatched_replay"], + counts["ubatched_eager"], + counts["ubatched_capture"], + counts["plain_eager"], + counts["plain_graph"], + ) + for key in counts: + counts[key] = 0 @property def graph_pool(self): @@ -153,25 +189,14 @@ def clear_graphs(self) -> None: @staticmethod def _create_sm_control_context(vllm_config: VllmConfig): - comm_sms: int = envs.VLLM_DBO_COMM_SMS + device = torch.accelerator.current_device_index() + total_sms = num_compute_units(device) + all2all_manager = get_all2all_manager_for_sm_control(vllm_config) + comm_sms = get_ubatch_comm_sms(vllm_config, total_sms, all2all_manager) set_comm_sms = lambda sms: None - if vllm_config.parallel_config.enable_expert_parallel: - # Currently only DeepEP highthroughput supports SM control so this - # only affects that case. - ep_group = get_ep_group() - device_communicator = ep_group.device_communicator - all2all_manager = None - if device_communicator is not None: - all2all_manager = device_communicator.all2all_manager - - if all2all_manager is not None: - max_sms_used = all2all_manager.max_sms_used() - if max_sms_used is not None: - comm_sms = min(comm_sms, max_sms_used) - - if comm_sms > 0 and all2all_manager is not None: - set_comm_sms = lambda sms: all2all_manager.set_num_sms(sms) + if comm_sms > 0 and all2all_manager is not None: + set_comm_sms = lambda sms: all2all_manager.set_num_sms(sms) # TODO(lucas): support other kernels besides DeepGEMM set_compute_sms = lambda sms: None @@ -343,6 +368,7 @@ def _make_ubatch_metadata( dp_metadata, batch_descriptor, cudagraph_runtime_mode, + for_capture: bool = False, ) -> list[UbatchMetadata]: # Create one forward context per ubatch forward_contexts = [] @@ -382,6 +408,7 @@ def _make_ubatch_metadata( positions, inputs_embeds, intermediate_tensors, + allow_clone=not for_capture, ) ubatch_metadata.append( UbatchMetadata( @@ -404,14 +431,23 @@ def _slice_model_inputs( positions, inputs_embeds, intermediate_tensors, + allow_clone: bool = True, ): - sliced_input_ids = input_ids[tokens_slice] if input_ids is not None else None + sliced_input_ids = ( + ensure_tensor_alignment(input_ids[tokens_slice], allow_clone=allow_clone) + if input_ids is not None + else None + ) # if we are using mrope. Mrope adds an additional dimension to the # positions tensor if positions.ndim == 2: sliced_positions = positions[:, tokens_slice] else: sliced_positions = positions[tokens_slice] + sliced_positions = ensure_tensor_alignment( + sliced_positions, + allow_clone=allow_clone, + ) sliced_inputs_embeds = ( inputs_embeds[tokens_slice] if inputs_embeds is not None else None ) @@ -447,11 +483,17 @@ def __call__(self, *args, **kwargs): if batch_descriptor.num_tokens in self.cudagraphs: cudagraph_runtime_mode = CUDAGraphMode.NONE - if cudagraph_runtime_mode in (CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE): - return self.runnable(*args, **kwargs) - else: - assert self.cudagraph_wrapper is not None - return self.cudagraph_wrapper(*args, **kwargs) + with self.sm_control: + if cudagraph_runtime_mode in ( + CUDAGraphMode.NONE, + CUDAGraphMode.PIECEWISE, + ): + self._count_mode("plain_eager") + return self.runnable(*args, **kwargs) + else: + assert self.cudagraph_wrapper is not None + self._count_mode("plain_graph") + return self.cudagraph_wrapper(*args, **kwargs) attn_metadata = forward_context.attn_metadata slot_mapping = forward_context.slot_mapping @@ -496,7 +538,9 @@ def __call__(self, *args, **kwargs): dp_metadata=ubatch_dp_metadata, batch_descriptor=batch_descriptor, cudagraph_runtime_mode=CUDAGraphMode.NONE, + for_capture=True, ) + self._count_mode("ubatched_capture") with self.sm_control: return self._capture_ubatches(ubatch_metadata, self.runnable) elif ( @@ -507,6 +551,7 @@ def __call__(self, *args, **kwargs): # Sync offloader before replay - ensures any external dependencies # from pre-capture prefetches are satisfied. get_offloader().sync_prev_onload() + self._count_mode("ubatched_replay") cudagraph_metadata.cudagraph.replay() return cudagraph_metadata.outputs else: @@ -523,5 +568,6 @@ def __call__(self, *args, **kwargs): batch_descriptor=batch_descriptor, cudagraph_runtime_mode=CUDAGraphMode.NONE, ) + self._count_mode("ubatched_eager") with self.sm_control: return self._run_ubatches(ubatch_metadata, self.runnable) diff --git a/vllm/v1/worker/sm_control.py b/vllm/v1/worker/sm_control.py new file mode 100644 index 000000000000..f4aec8b028a4 --- /dev/null +++ b/vllm/v1/worker/sm_control.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.utils.platform_utils import num_compute_units + + +_ZERO_COMM_SM_BACKENDS = frozenset({"deepep_low_latency", "nixl_ep"}) + + +def get_all2all_manager_for_sm_control(vllm_config: VllmConfig) -> Any | None: + if not vllm_config.parallel_config.enable_expert_parallel: + return None + + try: + from vllm.distributed import get_ep_group + + ep_group = get_ep_group() + except AssertionError: + return None + + device_communicator = ep_group.device_communicator + if device_communicator is None: + return None + return device_communicator.all2all_manager + + +def get_ubatch_comm_sms( + vllm_config: VllmConfig, + total_sms: int, + all2all_manager: Any | None = None, +) -> int: + if not vllm_config.parallel_config.use_ubatching: + return 0 + + comm_sms = envs.VLLM_DBO_COMM_SMS + if comm_sms <= 0: + return 0 + + if all2all_manager is None: + all2all_manager = get_all2all_manager_for_sm_control(vllm_config) + + if all2all_manager is not None: + max_sms_used = all2all_manager.max_sms_used() + if max_sms_used is not None: + comm_sms = min(comm_sms, max_sms_used) + elif vllm_config.parallel_config.all2all_backend in _ZERO_COMM_SM_BACKENDS: + comm_sms = 0 + + assert 0 <= comm_sms < total_sms, ( + f"Invalid ubatch comm SM count: comm_sms={comm_sms}, " + f"total_sms={total_sms}" + ) + return comm_sms + + +def get_ubatch_compute_sms( + vllm_config: VllmConfig, + device_index: int | None, + all2all_manager: Any | None = None, +) -> int: + total_sms = num_compute_units(device_index) + comm_sms = get_ubatch_comm_sms(vllm_config, total_sms, all2all_manager) + return total_sms - comm_sms diff --git a/vllm/v1/worker/ubatch_utils.py b/vllm/v1/worker/ubatch_utils.py index f4a76529023c..18a2ae871033 100644 --- a/vllm/v1/worker/ubatch_utils.py +++ b/vllm/v1/worker/ubatch_utils.py @@ -29,10 +29,32 @@ def num_tokens(self) -> int: UBatchSlices: TypeAlias = list[UBatchSlice] -def is_last_ubatch_empty( +def ensure_tensor_alignment( + tensor: torch.Tensor, alignment: int = 64, allow_clone: bool = True +) -> torch.Tensor: + if tensor.is_contiguous() and tensor.data_ptr() % alignment == 0: + return tensor + if not allow_clone: + # Clones made while preparing CUDA graph capture are baked into the + # captured kernels at their capture-time address. Replays would then + # silently consume stale data instead of the next batch's view. + raise RuntimeError( + f"ubatch slice is not contiguous/{alignment}B-aligned " + f"(shape={tuple(tensor.shape)}, dtype={tensor.dtype}, " + f"data_ptr%{alignment}={tensor.data_ptr() % alignment}) and " + "cloning is not allowed under CUDA graph capture; adjust " + "cudagraph_capture_sizes so ubatch splits stay aligned." + ) + return tensor.clone(memory_format=torch.contiguous_format) + + +def has_empty_ubatch( orig_num_tokens: int, padded_num_tokens: int, num_ubatches: int ) -> bool: - return (padded_num_tokens // num_ubatches) * (num_ubatches - 1) >= orig_num_tokens + tokens_per_ubatch = padded_num_tokens // num_ubatches + if tokens_per_ubatch == 0: + return True + return tokens_per_ubatch * (num_ubatches - 1) >= orig_num_tokens def check_ubatch_thresholds( @@ -51,6 +73,8 @@ def check_ubatch_thresholds( def _pad_out_ubatch_slices( ubatch_slices: UBatchSlices, num_total_tokens: int, num_reqs_padded: int ) -> UBatchSlices: + num_total_tokens = int(num_total_tokens) + num_reqs_padded = int(num_reqs_padded) last_slice = ubatch_slices[-1] padded_last_request_slice = slice(last_slice.request_slice.start, num_reqs_padded) padded_last_token_slice = slice(last_slice.token_slice.start, num_total_tokens) @@ -71,10 +95,18 @@ def maybe_create_ubatch_slices( if not should_ubatch: return None, None + num_tokens_padded = int(num_tokens_padded) + num_reqs_padded = int(num_reqs_padded) + num_ubatches = int(num_ubatches) + if split_point is None: - split_point = int(num_tokens_padded) // num_ubatches + split_point = num_tokens_padded // num_ubatches - token_split_points = [split_point * i for i in range(1, num_ubatches)] + if isinstance(split_point, list): + token_split_points = [int(point) for point in split_point] + else: + split_point = int(split_point) + token_split_points = [split_point * i for i in range(1, num_ubatches)] # TODO(lucas): Refactor the gpu_model_runner.py so we can pass # in cu_num_tokens directly (i.e. query_start_loc) @@ -85,9 +117,12 @@ def maybe_create_ubatch_slices( start_token = 0 # Add the end point to the split points to make iteration easier - all_points = token_split_points + [cu_num_tokens[-1]] + all_points = token_split_points + [int(cu_num_tokens[-1])] for end_token in all_points: + end_token = int(end_token) + if end_token <= start_token: + return None, None token_slice = slice(start_token, end_token) # Determine request slices using exclusive stop semantics @@ -103,7 +138,7 @@ def maybe_create_ubatch_slices( req_slice = slice(req_start, req_stop) ubatch_slices.append(UBatchSlice(req_slice, token_slice)) - start_token = end_token + start_token = int(end_token) ubatch_slices_padded = _pad_out_ubatch_slices( ubatch_slices, num_tokens_padded, num_reqs_padded @@ -132,7 +167,9 @@ def slice_query_start_locs( def _make_metadata_with_slice( - ubatch_slice: UBatchSlice, attn_metadata: CommonAttentionMetadata + ubatch_slice: UBatchSlice, + attn_metadata: CommonAttentionMetadata, + allow_clone: bool = True, ) -> CommonAttentionMetadata: """ This function creates a new CommonAttentionMetadata that corresponds to @@ -218,8 +255,8 @@ def _make_metadata_with_slice( # the attention backend selects the correct kernel for SWA layers. max_seq_len = max(int(seq_lens_cpu_upper_bound.max()), attn_metadata.max_seq_len) - num_requests = request_slice.stop - request_slice.start - num_actual_tokens = token_slice.stop - token_slice.start + num_requests = int(request_slice.stop) - int(request_slice.start) + num_actual_tokens = int(token_slice.stop) - int(token_slice.start) max_query_len = int( torch.max(torch.abs(query_start_loc_cpu[1:] - query_start_loc_cpu[:-1])).item() ) @@ -229,8 +266,27 @@ def _make_metadata_with_slice( if max_query_len == 0: max_query_len = attn_metadata.max_query_len - block_table_tensor = attn_metadata.block_table_tensor[request_slice] - slot_mapping = attn_metadata.slot_mapping[token_slice] + block_table_tensor = ensure_tensor_alignment( + attn_metadata.block_table_tensor[request_slice], + allow_clone=allow_clone, + ) + slot_mapping = ensure_tensor_alignment( + attn_metadata.slot_mapping[token_slice], + allow_clone=allow_clone, + ) + positions = ( + ensure_tensor_alignment( + attn_metadata.positions[token_slice], + allow_clone=allow_clone, + ) + if attn_metadata.positions is not None + else None + ) + is_prefilling = ( + attn_metadata.is_prefilling[request_slice] + if attn_metadata.is_prefilling is not None + else None + ) return CommonAttentionMetadata( query_start_loc=query_start_loc, @@ -242,6 +298,8 @@ def _make_metadata_with_slice( max_seq_len=max_seq_len, block_table_tensor=block_table_tensor, slot_mapping=slot_mapping, + positions=positions, + is_prefilling=is_prefilling, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, _seq_lens_cpu=seq_lens_cpu, _num_computed_tokens_cpu=num_computed_tokens_cpu, @@ -251,15 +309,23 @@ def _make_metadata_with_slice( def split_attn_metadata( ubatch_slices: list[UBatchSlice], common_attn_metadata: CommonAttentionMetadata, + allow_clone: bool = True, ) -> list[CommonAttentionMetadata]: """ Creates a new CommonAttentionMetadata instance that corresponds to the requests for each UBatchSlice in ubatch_slices. Note: This function does not modify common_attn_metadata + Pass allow_clone=False when the metadata feeds CUDA graph capture. """ results = [] for ubatch_slice in ubatch_slices: - results.append(_make_metadata_with_slice(ubatch_slice, common_attn_metadata)) + results.append( + _make_metadata_with_slice( + ubatch_slice, + common_attn_metadata, + allow_clone=allow_clone, + ) + ) return results