| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Callable |
| 4 | from dataclasses import dataclass |
| 5 | from typing import TypeAlias |
| 6 | |
| 7 | import numpy as np |
| 8 | import torch |
| 9 | |
| 10 | import vllm.envs as envs |
| 11 | from vllm.config import ParallelConfig |
| 12 | from vllm.distributed import get_ep_group |
| 13 | from vllm.platforms import current_platform |
| 14 | from vllm.utils.deep_gemm import set_num_sms as deep_gemm_set_num_sms |
| 15 | from vllm.utils.import_utils import has_deep_gemm |
| 16 | from vllm.utils.platform_utils import num_compute_units |
| 17 | from vllm.v1.attention.backend import CommonAttentionMetadata |
| 18 | |
| 19 | |
| 20 | @dataclass |
| 21 | class UBatchSlice: |
| 22 | request_slice: slice |
| 23 | token_slice: slice |
| 24 | |
| 25 | def is_empty(self) -> bool: |
| 26 | return ( |
| 27 | self.request_slice.start == self.request_slice.stop |
| 28 | or self.token_slice.start == self.token_slice.stop |
| 29 | ) |
| 30 | |
| 31 | @property |
| 32 | def num_tokens(self) -> int: |
| 33 | return self.token_slice.stop - self.token_slice.start |
| 34 | |
| 35 | |
| 36 | UBatchSlices: TypeAlias = list[UBatchSlice] |
| 37 | |
| 38 | |
| 39 | class SMControlContextManager: |
| 40 | def __init__( |
| 41 | self, |
| 42 | comm_sms: int, |
| 43 | set_comm_sms: Callable[[int], None], |
| 44 | set_compute_sms: Callable[[int], None], |
| 45 | ): |
| 46 | """Context manager for controlling SM (Streaming Multiprocessor) |
| 47 | allocation. Upon entering the context, it sets the number of SMs |
| 48 | allocated for communication and computation to comm_sms and |
| 49 | total_sms - comm_sms respectively. Upon exiting, it restores the |
| 50 | allocation to use all available SMs (i.e. total_sms). |
| 51 | |
| 52 | Args: |
| 53 | comm_sms (int): The number of SMs to allocate for communication. |
| 54 | (The remainder will be used for computation.) |
| 55 | set_comm_sms (Callable[[int], None]): |
| 56 | A function that sets the number of SMs for communication. |
| 57 | set_compute_sms (Callable[[int], None]): |
| 58 | A function that sets the number of SMs for computation. |
| 59 | |
| 60 | """ |
| 61 | assert current_platform.is_cuda() or current_platform.is_rocm(), ( |
| 62 | "SM/CU control is supported on CUDA and ROCm platforms" |
| 63 | ) |
| 64 | device = torch.accelerator.current_device_index() |
| 65 | total_sms = num_compute_units(device) |
| 66 | |
| 67 | assert comm_sms < total_sms |
| 68 | self.total_sms = total_sms |
| 69 | self.compute_sms = total_sms - comm_sms |
| 70 | self.comm_sms = comm_sms |
| 71 | self.set_comm_sms = set_comm_sms |
| 72 | self.set_compute_sms = set_compute_sms |
| 73 | |
| 74 | def __enter__(self): |
| 75 | self.set_comm_sms(self.comm_sms) |
| 76 | self.set_compute_sms(self.compute_sms) |
| 77 | |
| 78 | def __exit__(self, exc_type, exc_value, traceback): |
| 79 | self.set_comm_sms(self.total_sms) |
| 80 | self.set_compute_sms(self.total_sms) |
| 81 | |
| 82 | |
| 83 | def create_sm_control_context( |
| 84 | parallel_config: ParallelConfig, |
| 85 | ) -> SMControlContextManager: |
| 86 | """Reserve SMs for communication kernels while microbatches overlap.""" |
| 87 | comm_sms: int = envs.VLLM_DBO_COMM_SMS |
| 88 | rocm_deepep_ht_dbo = ( |
| 89 | current_platform.is_rocm() |
| 90 | and parallel_config.enable_dbo |
| 91 | and parallel_config.all2all_backend == "deepep_high_throughput" |
| 92 | ) |
| 93 | if rocm_deepep_ht_dbo: |
| 94 | # On ROCm, reserving CUs for DeepEP HT communication under DBO |
| 95 | # corrupts DP+EP generation accuracy. Keep the backend active, but |
| 96 | # leave all CUs visible to the compute and communication kernels. |
| 97 | comm_sms = 0 |
| 98 | |
| 99 | set_comm_sms = lambda sms: None |
| 100 | if parallel_config.enable_expert_parallel: |
| 101 | # Currently only DeepEP highthroughput supports SM control so this |
| 102 | # only affects that case. |
| 103 | ep_group = get_ep_group() |
| 104 | device_communicator = ep_group.device_communicator |
| 105 | all2all_manager = None |
| 106 | if device_communicator is not None: |
| 107 | all2all_manager = device_communicator.all2all_manager |
| 108 | |
| 109 | if all2all_manager is not None: |
| 110 | max_sms_used = all2all_manager.max_sms_used() |
| 111 | if max_sms_used is not None: |
| 112 | comm_sms = min(comm_sms, max_sms_used) |
| 113 | |
| 114 | if comm_sms > 0 and all2all_manager is not None: |
| 115 | set_comm_sms = lambda sms: all2all_manager.set_num_sms(sms) |
| 116 | |
| 117 | # TODO(lucas): support other kernels besides DeepGEMM |
| 118 | set_compute_sms = lambda sms: None |
| 119 | if has_deep_gemm() and comm_sms > 0: |
| 120 | set_compute_sms = lambda sms: deep_gemm_set_num_sms(sms) |
| 121 | |
| 122 | return SMControlContextManager( |
| 123 | comm_sms=comm_sms, |
| 124 | set_comm_sms=set_comm_sms, |
| 125 | set_compute_sms=set_compute_sms, |
| 126 | ) |
| 127 | |
| 128 | |
| 129 | def is_last_ubatch_empty( |
| 130 | orig_num_tokens: int, padded_num_tokens: int, num_ubatches: int |
| 131 | ) -> bool: |
| 132 | return (padded_num_tokens // num_ubatches) * (num_ubatches - 1) >= orig_num_tokens |
| 133 | |
| 134 | |
| 135 | def get_num_ubatches(parallel_config: ParallelConfig) -> int: |
| 136 | """How many microbatches a step is split into; 1 when microbatching is off.""" |
| 137 | return parallel_config.num_ubatches if parallel_config.use_ubatching else 1 |
| 138 | |
| 139 | |
| 140 | def check_ubatch_thresholds( |
| 141 | config: ParallelConfig, num_tokens: int, uniform_decode: bool |
| 142 | ) -> bool: |
| 143 | if not config.use_ubatching: |
| 144 | return False |
| 145 | if uniform_decode: |
| 146 | return num_tokens >= config.dbo_decode_token_threshold |
| 147 | else: |
| 148 | return num_tokens >= config.dbo_prefill_token_threshold |
| 149 | |
| 150 | |
| 151 | # This pads the last ubatch slice out to the total number of tokens |
| 152 | # (num_tokens + padding) since we do `create_ubatch_slices` before applying DP padding. |
| 153 | def _pad_out_ubatch_slices( |
| 154 | ubatch_slices: UBatchSlices, num_total_tokens: int, num_reqs_padded: int |
| 155 | ) -> UBatchSlices: |
| 156 | last_slice = ubatch_slices[-1] |
| 157 | padded_last_request_slice = slice(last_slice.request_slice.start, num_reqs_padded) |
| 158 | padded_last_token_slice = slice(last_slice.token_slice.start, num_total_tokens) |
| 159 | |
| 160 | return ubatch_slices[:-1] + [ |
| 161 | UBatchSlice(padded_last_request_slice, padded_last_token_slice) |
| 162 | ] |
| 163 | |
| 164 | |
| 165 | def maybe_create_ubatch_slices( |
| 166 | should_ubatch: bool, |
| 167 | num_scheduled_tokens: np.ndarray, |
| 168 | num_tokens_padded: int, |
| 169 | num_reqs_padded: int, |
| 170 | num_ubatches: int, |
| 171 | split_point: list[int] | int | None = None, |
| 172 | ) -> tuple[UBatchSlices | None, UBatchSlices | None]: |
| 173 | if not should_ubatch: |
| 174 | return None, None |
| 175 | |
| 176 | if split_point is None: |
| 177 | split_point = int(num_tokens_padded) // num_ubatches |
| 178 | |
| 179 | token_split_points = [split_point * i for i in range(1, num_ubatches)] |
| 180 | |
| 181 | # TODO(lucas): Refactor the gpu_model_runner.py so we can pass |
| 182 | # in cu_num_tokens directly (i.e. query_start_loc) |
| 183 | cu_num_tokens = np.zeros(len(num_scheduled_tokens) + 1, dtype=np.int32) |
| 184 | np.cumsum(num_scheduled_tokens, dtype=np.int32, out=cu_num_tokens[1:]) |
| 185 | |
| 186 | ubatch_slices = [] |
| 187 | start_token = 0 |
| 188 | |
| 189 | # Add the end point to the split points to make iteration easier |
| 190 | all_points = token_split_points + [cu_num_tokens[-1]] |
| 191 | |
| 192 | for end_token in all_points: |
| 193 | token_slice = slice(start_token, end_token) |
| 194 | |
| 195 | # Determine request slices using exclusive stop semantics |
| 196 | # Ubatch includes requests whose tokens overlap [start_token, end_token) |
| 197 | |
| 198 | # Start at the request that contains the start_token |
| 199 | # or the request starting exactly at start_token (if on boundary) |
| 200 | req_start = int(np.searchsorted(cu_num_tokens, start_token, side="right") - 1) |
| 201 | |
| 202 | # Stop at the request that starts at or after end_token |
| 203 | req_stop = int(np.searchsorted(cu_num_tokens, end_token, side="left")) |
| 204 | |
| 205 | req_slice = slice(req_start, req_stop) |
| 206 | ubatch_slices.append(UBatchSlice(req_slice, token_slice)) |
| 207 | |
| 208 | start_token = end_token |
| 209 | |
| 210 | ubatch_slices_padded = _pad_out_ubatch_slices( |
| 211 | ubatch_slices, num_tokens_padded, num_reqs_padded |
| 212 | ) |
| 213 | |
| 214 | assert sum(s.num_tokens for s in ubatch_slices_padded) == num_tokens_padded |
| 215 | |
| 216 | return ubatch_slices, ubatch_slices_padded |
| 217 | |
| 218 | |
| 219 | def slice_query_start_locs( |
| 220 | query_start_loc: torch.Tensor, |
| 221 | request_slice: slice, |
| 222 | ) -> torch.Tensor: |
| 223 | """Creates a new query_start_loc that corresponds to the requests in |
| 224 | request_slice. |
| 225 | |
| 226 | Note: This function creates a new tensor to hold the new query_start_locs. |
| 227 | This will break cudagraph compatibility. |
| 228 | """ |
| 229 | return ( |
| 230 | query_start_loc[request_slice.start : request_slice.stop + 1] |
| 231 | - query_start_loc[request_slice.start] |
| 232 | ) |
| 233 | |
| 234 | |
| 235 | def _make_metadata_with_slice( |
| 236 | ubatch_slice: UBatchSlice, attn_metadata: CommonAttentionMetadata |
| 237 | ) -> CommonAttentionMetadata: |
| 238 | """This function creates a new CommonAttentionMetadata that corresponds to |
| 239 | the requests included in ubatch_slice |
| 240 | """ |
| 241 | assert not ubatch_slice.is_empty(), f"Ubatch slice {ubatch_slice} is empty" |
| 242 | |
| 243 | request_slice = ubatch_slice.request_slice |
| 244 | token_slice = ubatch_slice.token_slice |
| 245 | |
| 246 | start_locs = attn_metadata.query_start_loc_cpu |
| 247 | first_req = request_slice.start |
| 248 | first_tok = token_slice.start |
| 249 | last_req = request_slice.stop - 1 |
| 250 | last_tok = token_slice.stop - 1 |
| 251 | |
| 252 | assert start_locs[first_req] <= first_tok < start_locs[first_req + 1], ( |
| 253 | "Token slice start outside of first request" |
| 254 | ) |
| 255 | # NOTE: last token can be outside of the last request if we have CG padding. |
| 256 | |
| 257 | # If the request is split across ubatches, we have to adjust the metadata. |
| 258 | # splits_first_request: The first request in this slice is the continuation of |
| 259 | # a request that started in a previous slice. |
| 260 | # splits_last_request: The last request in this slice continues into the |
| 261 | # next slice. |
| 262 | splits_first_request = first_tok > start_locs[first_req] |
| 263 | splits_last_request = last_tok < start_locs[last_req + 1] - 1 |
| 264 | |
| 265 | query_start_loc_cpu = slice_query_start_locs(start_locs, request_slice) |
| 266 | query_start_loc = slice_query_start_locs( |
| 267 | attn_metadata.query_start_loc, request_slice |
| 268 | ) |
| 269 | |
| 270 | assert len(query_start_loc) >= 2, ( |
| 271 | f"query_start_loc must have at least 2 elements, got {len(query_start_loc)}" |
| 272 | ) |
| 273 | |
| 274 | if splits_first_request: |
| 275 | tokens_skipped = first_tok - start_locs[first_req] |
| 276 | query_start_loc[1:] -= tokens_skipped |
| 277 | query_start_loc_cpu[1:] -= tokens_skipped |
| 278 | seq_lens = attn_metadata.seq_lens[request_slice] |
| 279 | seq_lens_cpu_upper_bound = ( |
| 280 | attn_metadata.seq_lens_cpu_upper_bound[request_slice] |
| 281 | if attn_metadata.seq_lens_cpu_upper_bound is not None |
| 282 | else None |
| 283 | ) |
| 284 | if splits_last_request: |
| 285 | # NOTE: We use start_locs (the original query_start_loc_cpu) to calculate |
| 286 | # the tokens skipped because query_start_loc_cpu might have been modified |
| 287 | # if splits_first_request is True. |
| 288 | tokens_skipped = start_locs[last_req + 1] - token_slice.stop |
| 289 | query_start_loc[-1] -= tokens_skipped |
| 290 | query_start_loc_cpu[-1] -= tokens_skipped |
| 291 | |
| 292 | # Make sure we don't modify the seq_lens tensors |
| 293 | # (not cudagraph compatible) |
| 294 | seq_lens = seq_lens.clone() |
| 295 | seq_lens[-1] -= tokens_skipped |
| 296 | if seq_lens_cpu_upper_bound is not None: |
| 297 | seq_lens_cpu_upper_bound = seq_lens_cpu_upper_bound.clone() |
| 298 | seq_lens_cpu_upper_bound[-1] -= tokens_skipped |
| 299 | |
| 300 | assert seq_lens_cpu_upper_bound is not None |
| 301 | # Preserve the max_seq_len override set during CUDA-graph capture so |
| 302 | # the attention backend selects the correct kernel for SWA layers. |
| 303 | max_seq_len = max(int(seq_lens_cpu_upper_bound.max()), attn_metadata.max_seq_len) |
| 304 | |
| 305 | num_requests = request_slice.stop - request_slice.start |
| 306 | num_actual_tokens = token_slice.stop - token_slice.start |
| 307 | max_query_len = int( |
| 308 | torch.max(torch.abs(query_start_loc_cpu[1:] - query_start_loc_cpu[:-1])).item() |
| 309 | ) |
| 310 | |
| 311 | # This is to account for the case where we are in a dummy |
| 312 | # run and query_start_loc_cpu is full of 0s |
| 313 | if max_query_len == 0: |
| 314 | max_query_len = attn_metadata.max_query_len |
| 315 | |
| 316 | block_table_tensor = attn_metadata.block_table_tensor[request_slice] |
| 317 | slot_mapping = attn_metadata.slot_mapping[token_slice] |
| 318 | |
| 319 | return CommonAttentionMetadata( |
| 320 | query_start_loc=query_start_loc, |
| 321 | query_start_loc_cpu=query_start_loc_cpu, |
| 322 | seq_lens=seq_lens, |
| 323 | num_reqs=num_requests, |
| 324 | num_actual_tokens=num_actual_tokens, |
| 325 | max_query_len=max_query_len, |
| 326 | max_seq_len=max_seq_len, |
| 327 | block_table_tensor=block_table_tensor, |
| 328 | slot_mapping=slot_mapping, |
| 329 | seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, |
| 330 | ) |
| 331 | |
| 332 | |
| 333 | def split_attn_metadata( |
| 334 | ubatch_slices: list[UBatchSlice], |
| 335 | common_attn_metadata: CommonAttentionMetadata, |
| 336 | ) -> list[CommonAttentionMetadata]: |
| 337 | """Creates a new CommonAttentionMetadata instance that corresponds to the |
| 338 | requests for each UBatchSlice in ubatch_slices. |
| 339 | |
| 340 | Note: This function does not modify common_attn_metadata |
| 341 | """ |
| 342 | results = [] |
| 343 | for ubatch_slice in ubatch_slices: |
| 344 | results.append(_make_metadata_with_slice(ubatch_slice, common_attn_metadata)) |
| 345 | |
| 346 | return results |