| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import replace |
| 4 | from typing import Any |
| 5 | |
| 6 | import numpy as np |
| 7 | import torch |
| 8 | import torch.distributed as dist |
| 9 | import torch.nn as nn |
| 10 | |
| 11 | from vllm.config import CUDAGraphMode, VllmConfig |
| 12 | from vllm.distributed.parallel_state import get_dp_group |
| 13 | from vllm.forward_context import DPMetadata, create_forward_context |
| 14 | from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers |
| 15 | from vllm.triton_utils import tl, triton |
| 16 | from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadataBuilder |
| 17 | from vllm.v1.attention.backends.utils import PAD_SLOT_ID |
| 18 | from vllm.v1.core.sched.output import NewRequestData |
| 19 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 20 | from vllm.v1.worker.dp_utils import should_skip_dp_coordination |
| 21 | from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer |
| 22 | from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool |
| 23 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 24 | from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache |
| 25 | from vllm.v1.worker.gpu.model_states.default import DefaultModelState |
| 26 | from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata |
| 27 | from vllm.v1.worker.gpu.states import RequestState |
| 28 | from vllm.v1.worker.utils import AttentionGroup |
| 29 | |
| 30 | |
| 31 | @triton.jit |
| 32 | def _gather_lookback_kernel( |
| 33 | lookback_ptr, |
| 34 | idx_mapping_ptr, |
| 35 | num_computed_tokens_ptr, |
| 36 | all_token_ids_ptr, |
| 37 | all_token_ids_stride, |
| 38 | num_reqs, |
| 39 | DEPTH: tl.constexpr, |
| 40 | BLOCK_DEPTH: tl.constexpr, |
| 41 | ): |
| 42 | # One program per lookback row; rows past the batch are filled with -1. |
| 43 | batch_idx = tl.program_id(0) |
| 44 | in_batch = batch_idx < num_reqs |
| 45 | req_state_idx = tl.load(idx_mapping_ptr + batch_idx, mask=in_batch, other=0) |
| 46 | num_computed = tl.load(num_computed_tokens_ptr + req_state_idx) |
| 47 | |
| 48 | offs = tl.arange(0, BLOCK_DEPTH) |
| 49 | pos = num_computed - 1 - offs |
| 50 | valid = in_batch & (offs < DEPTH) & (pos >= 0) |
| 51 | ids = tl.load( |
| 52 | all_token_ids_ptr + req_state_idx * all_token_ids_stride + pos, |
| 53 | mask=valid, |
| 54 | other=-1, |
| 55 | ) |
| 56 | tl.store(lookback_ptr + batch_idx * DEPTH + offs, ids, mask=offs < DEPTH) |
| 57 | |
| 58 | |
| 59 | @triton.jit |
| 60 | def _pad_replayed_slots_kernel( |
| 61 | slot_mappings_ptr, # [num_groups, num_tokens_padded] |
| 62 | group_stride, |
| 63 | cacheable_groups_ptr, # [NUM_GROUPS] indices of the prefix-cacheable groups |
| 64 | query_start_loc_ptr, # [num_reqs + 1] |
| 65 | positions_ptr, # [num_tokens] |
| 66 | replay_start_ptr, # [num_reqs] |
| 67 | window, |
| 68 | pad_slot_id, |
| 69 | NUM_GROUPS: tl.constexpr, |
| 70 | BLOCK: tl.constexpr, |
| 71 | ): |
| 72 | req = tl.program_id(0) |
| 73 | start = tl.load(replay_start_ptr + req) |
| 74 | if start == 0: |
| 75 | return |
| 76 | begin = tl.load(query_start_loc_ptr + req) |
| 77 | end = tl.load(query_start_loc_ptr + req + 1) |
| 78 | for tok in range(begin, end, BLOCK): |
| 79 | offs = tok + tl.arange(0, BLOCK) |
| 80 | pos = tl.load(positions_ptr + offs, mask=offs < end, other=0) |
| 81 | replayed = (offs < end) & (pos >= start) & (pos < start + window) |
| 82 | for i in tl.static_range(NUM_GROUPS): |
| 83 | group = tl.load(cacheable_groups_ptr + i) |
| 84 | tl.store( |
| 85 | slot_mappings_ptr + group * group_stride + offs, |
| 86 | pad_slot_id, |
| 87 | mask=replayed, |
| 88 | ) |
| 89 | |
| 90 | |
| 91 | @triton.jit |
| 92 | def _gather_replay_batch_kernel( |
| 93 | query_start_loc_ptr, # [num_reqs + 1] |
| 94 | dropped_before_ptr, # [num_reqs + 1] rows trimmed before each boundary |
| 95 | seq_lens_ptr, # [num_reqs] |
| 96 | replay_start_ptr, # [num_reqs] the encoder-side replay start |
| 97 | positions_ptr, |
| 98 | slot_mappings_ptr, # [num_groups, num_tokens] |
| 99 | slot_mappings_stride, |
| 100 | rows_ptr, # out: [num_kept] batch row of every kept row |
| 101 | kept_query_start_loc_ptr, # out: [num_reqs + 1] |
| 102 | kept_positions_ptr, # out: [num_kept] |
| 103 | kept_slot_mappings_ptr, # out: [num_groups, num_kept] |
| 104 | kept_slot_mappings_stride, |
| 105 | kept_kv_start_ptr, # out: [num_reqs] |
| 106 | window, |
| 107 | NUM_GROUPS: tl.constexpr, |
| 108 | BLOCK: tl.constexpr, |
| 109 | ): |
| 110 | """One program per request: its kept rows are the last ones before its |
| 111 | boundary, and land at that boundary minus the rows trimmed before it.""" |
| 112 | req = tl.program_id(0) |
| 113 | begin = tl.load(query_start_loc_ptr + req) |
| 114 | end = tl.load(query_start_loc_ptr + req + 1) |
| 115 | kept_begin = begin - tl.load(dropped_before_ptr + req) |
| 116 | kept_end = end - tl.load(dropped_before_ptr + req + 1) |
| 117 | if req == 0: |
| 118 | tl.store(kept_query_start_loc_ptr, kept_begin) |
| 119 | tl.store(kept_query_start_loc_ptr + req + 1, kept_end) |
| 120 | |
| 121 | # A trimmed request holds no replay-layer window KV below its kept rows. |
| 122 | kv_start = tl.load(replay_start_ptr + req) |
| 123 | if kept_end - kept_begin < end - begin: |
| 124 | kv_start = tl.maximum(kv_start, tl.load(seq_lens_ptr + req) - window) |
| 125 | tl.store(kept_kv_start_ptr + req, kv_start) |
| 126 | |
| 127 | first_row = end - (kept_end - kept_begin) |
| 128 | for tok in range(kept_begin, kept_end, BLOCK): |
| 129 | offs = tok + tl.arange(0, BLOCK) |
| 130 | mask = offs < kept_end |
| 131 | row = first_row + (offs - kept_begin) |
| 132 | tl.store(rows_ptr + offs, row.to(tl.int64), mask=mask) |
| 133 | tl.store( |
| 134 | kept_positions_ptr + offs, |
| 135 | tl.load(positions_ptr + row, mask=mask), |
| 136 | mask=mask, |
| 137 | ) |
| 138 | for g in tl.static_range(NUM_GROUPS): |
| 139 | tl.store( |
| 140 | kept_slot_mappings_ptr + g * kept_slot_mappings_stride + offs, |
| 141 | tl.load(slot_mappings_ptr + g * slot_mappings_stride + row, mask=mask), |
| 142 | mask=mask, |
| 143 | ) |
| 144 | |
| 145 | |
| 146 | class ReplayAttnMetadata(ModelSpecificAttnMetadata): |
| 147 | """Hands the batch's replay starts to the sliding-window builders.""" |
| 148 | |
| 149 | def __init__(self, replay_start: torch.Tensor) -> None: |
| 150 | self.replay_start = replay_start |
| 151 | |
| 152 | def get_extra_attn_kwargs( |
| 153 | self, attn_metadata_builder: Any, num_reqs: int |
| 154 | ) -> dict[str, Any]: |
| 155 | if isinstance(attn_metadata_builder, DeepseekSparseSWAMetadataBuilder): |
| 156 | return {"replay_start": self.replay_start[:num_reqs]} |
| 157 | return {} |
| 158 | |
| 159 | |
| 160 | class DeepseekV41ModelState(DefaultModelState): |
| 161 | """DefaultModelState plus the engram lookback window and SWA bounded replay. |
| 162 | |
| 163 | The engram n-gram hash needs the ids of the ``depth`` tokens preceding |
| 164 | each request's chunk start (see ``common/engram.py``). The runner keeps |
| 165 | the full token history on device, so the window is gathered there every |
| 166 | step: exact for prompt and generated tokens alike, whatever instance |
| 167 | produced their KV. |
| 168 | |
| 169 | SWA bounded replay (``CacheConfig.swa_bounded_replay``) keeps the |
| 170 | sliding-window KV out of prefix caching and rebuilds it after a prefix hit |
| 171 | by recomputing the hit's last window; the scheduler tells each request the |
| 172 | position from which it holds window KV (``NewRequestData.replay_start``). |
| 173 | ``prepare_attn`` gathers those per batch, pads the replayed tokens' slots |
| 174 | in the prefix-cacheable groups so the cached KV stays as is, and hands the |
| 175 | starts to the sliding-window metadata builders, whose kernels read no |
| 176 | window KV below them. With the decoder side on |
| 177 | (``model.decoder_replay_layers``) it also prepares the replay layers' |
| 178 | batch in eager steps: each prefill's last ``window`` rows as a sub-batch |
| 179 | with metadata and a forward context of its own, like a microbatch. |
| 180 | """ |
| 181 | |
| 182 | def __init__( |
| 183 | self, |
| 184 | vllm_config: VllmConfig, |
| 185 | model: nn.Module, |
| 186 | encoder_cache: EncoderCache | None, |
| 187 | device: torch.device, |
| 188 | ): |
| 189 | super().__init__(vllm_config, model, encoder_cache, device) |
| 190 | depth = model.token_lookback_depth |
| 191 | self.lookback_token_ids: torch.Tensor | None = None |
| 192 | if depth > 0: |
| 193 | # Persistent so a captured graph can read it on replay. |
| 194 | self.lookback_token_ids = torch.full( |
| 195 | (self.max_num_reqs, depth), -1, dtype=torch.int32, device=device |
| 196 | ) |
| 197 | |
| 198 | # Per request state index; batches gather from it (see prepare_attn). |
| 199 | self._replay_start_np = np.zeros(self.max_num_reqs, dtype=np.int32) |
| 200 | self._replay_start = torch.zeros( |
| 201 | self.max_num_reqs, dtype=torch.int32, device=device |
| 202 | ) |
| 203 | self._replay_start_staging = UvaBufferPool(self.max_num_reqs, torch.int32) |
| 204 | # The replay window and the indices of the prefix-cacheable groups, from |
| 205 | # the KV cache config on first use; None until then. |
| 206 | self._replay: tuple[int, torch.Tensor] | None = None |
| 207 | |
| 208 | self.decoder_replay_layers: DecoderReplayLayers | None = getattr( |
| 209 | model, "decoder_replay_layers", None |
| 210 | ) |
| 211 | if self.decoder_replay_layers is not None: |
| 212 | # Requests whose every prompt row is read (prompt logprobs) never trim. |
| 213 | self._keeps_rows_np = np.zeros(self.max_num_reqs, dtype=np.bool_) |
| 214 | self._dropped_before = UvaBufferPool(self.max_num_reqs + 1, torch.int32) |
| 215 | # The replay batch's inputs: its rows of the batch, query offsets, |
| 216 | # positions, slot mappings (sized with the KV cache groups on first |
| 217 | # use) and the position from which each request holds replay-layer |
| 218 | # window KV. |
| 219 | self._kept_rows = torch.zeros( |
| 220 | self.max_num_tokens, dtype=torch.int64, device=device |
| 221 | ) |
| 222 | self._kept_query_start_loc = torch.zeros( |
| 223 | self.max_num_reqs + 1, dtype=torch.int32, device=device |
| 224 | ) |
| 225 | self._kept_positions = torch.zeros( |
| 226 | self.max_num_tokens, dtype=torch.int64, device=device |
| 227 | ) |
| 228 | self._kept_slot_mappings: torch.Tensor | None = None |
| 229 | self._kept_kv_start = torch.zeros( |
| 230 | self.max_num_reqs, dtype=torch.int32, device=device |
| 231 | ) |
| 232 | self._replay_attn_groups: list[list[AttentionGroup]] | None = None |
| 233 | |
| 234 | def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: |
| 235 | super().add_request(req_index, new_req_data) |
| 236 | self._replay_start_np[req_index] = new_req_data.replay_start |
| 237 | if self.decoder_replay_layers is not None: |
| 238 | params = new_req_data.sampling_params |
| 239 | self._keeps_rows_np[req_index] = ( |
| 240 | params is not None and params.prompt_logprobs is not None |
| 241 | ) |
| 242 | |
| 243 | def prepare_inputs( |
| 244 | self, input_batch: InputBatch, req_states: RequestState |
| 245 | ) -> dict[str, torch.Tensor | None]: |
| 246 | model_inputs = super().prepare_inputs(input_batch, req_states) |
| 247 | window = self.lookback_token_ids |
| 248 | if window is None: |
| 249 | return model_inputs |
| 250 | all_token_ids = req_states.all_token_ids.gpu |
| 251 | depth = window.shape[1] |
| 252 | _gather_lookback_kernel[(window.shape[0],)]( |
| 253 | window, |
| 254 | input_batch.idx_mapping, |
| 255 | req_states.num_computed_tokens.gpu, |
| 256 | all_token_ids, |
| 257 | all_token_ids.stride(0), |
| 258 | input_batch.idx_mapping.shape[0], |
| 259 | DEPTH=depth, |
| 260 | BLOCK_DEPTH=triton.next_power_of_2(depth), |
| 261 | ) |
| 262 | model_inputs["lookback_token_ids"] = window |
| 263 | return model_inputs |
| 264 | |
| 265 | def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]: |
| 266 | model_inputs = super().prepare_dummy_inputs(num_reqs, num_tokens) |
| 267 | if self.lookback_token_ids is not None: |
| 268 | # The captured graph reads this buffer; replays refill it in place. |
| 269 | self.lookback_token_ids.fill_(-1) |
| 270 | model_inputs["lookback_token_ids"] = self.lookback_token_ids |
| 271 | return model_inputs |
| 272 | |
| 273 | def prepare_attn( |
| 274 | self, |
| 275 | input_batch: InputBatch, |
| 276 | cudagraph_mode: CUDAGraphMode, |
| 277 | block_tables: tuple[torch.Tensor, ...], |
| 278 | slot_mappings: torch.Tensor, |
| 279 | attn_groups: list[list[AttentionGroup]], |
| 280 | kv_cache_config: KVCacheConfig, |
| 281 | for_capture: bool = False, |
| 282 | ubatch_idx: int = 0, |
| 283 | model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, |
| 284 | ) -> dict[str, Any]: |
| 285 | if self._replay is None: |
| 286 | specs = [group.kv_cache_spec for group in kv_cache_config.kv_cache_groups] |
| 287 | self._replay = ( |
| 288 | max(spec.prefix_replay_tokens for spec in specs), |
| 289 | torch.tensor( |
| 290 | [i for i, spec in enumerate(specs) if spec.prefix_cacheable], |
| 291 | dtype=torch.int32, |
| 292 | device=self.device, |
| 293 | ), |
| 294 | ) |
| 295 | if self.decoder_replay_layers is not None: |
| 296 | self._kept_slot_mappings = torch.zeros( |
| 297 | len(specs), |
| 298 | self.max_num_tokens, |
| 299 | dtype=torch.int64, |
| 300 | device=self.device, |
| 301 | ) |
| 302 | self._warm_up_replay_kernels(input_batch, slot_mappings) |
| 303 | window, cacheable_groups = self._replay |
| 304 | replay_start: torch.Tensor | None = None |
| 305 | if window: |
| 306 | num_reqs = input_batch.num_reqs |
| 307 | # Decode rows sit above the hit, so only prefills carry a replay |
| 308 | # start; dummy batches (captures, profiling) carry none. |
| 309 | replay_start_np = np.where( |
| 310 | input_batch.is_prefilling_np[:num_reqs], |
| 311 | self._replay_start_np[input_batch.idx_mapping_np[:num_reqs]], |
| 312 | 0, |
| 313 | ).astype(np.int32) |
| 314 | replay_start = self._replay_start_staging.copy_to_gpu( |
| 315 | replay_start_np, out=self._replay_start[:num_reqs] |
| 316 | ) |
| 317 | if replay_start_np.any(): |
| 318 | # The replayed tokens rebuild window KV only: their slots in the |
| 319 | # prefix-cacheable groups are padded so the cached KV stays as is. |
| 320 | _pad_replayed_slots_kernel[(num_reqs,)]( |
| 321 | slot_mappings, |
| 322 | slot_mappings.stride(0), |
| 323 | cacheable_groups, |
| 324 | input_batch.query_start_loc, |
| 325 | input_batch.positions, |
| 326 | replay_start, |
| 327 | window, |
| 328 | PAD_SLOT_ID, |
| 329 | NUM_GROUPS=cacheable_groups.numel(), |
| 330 | BLOCK=1024, |
| 331 | ) |
| 332 | assert model_specific_attn_metadata is None |
| 333 | model_specific_attn_metadata = ReplayAttnMetadata(replay_start) |
| 334 | attn_metadata = super().prepare_attn( |
| 335 | input_batch, |
| 336 | cudagraph_mode, |
| 337 | block_tables, |
| 338 | slot_mappings, |
| 339 | attn_groups, |
| 340 | kv_cache_config, |
| 341 | for_capture=for_capture, |
| 342 | ubatch_idx=ubatch_idx, |
| 343 | model_specific_attn_metadata=model_specific_attn_metadata, |
| 344 | ) |
| 345 | if self.decoder_replay_layers is not None: |
| 346 | assert replay_start is not None |
| 347 | self._prepare_replay_batch( |
| 348 | input_batch, |
| 349 | cudagraph_mode, |
| 350 | block_tables, |
| 351 | slot_mappings, |
| 352 | attn_groups, |
| 353 | kv_cache_config, |
| 354 | replay_start, |
| 355 | ) |
| 356 | return attn_metadata |
| 357 | |
| 358 | def _warm_up_replay_kernels( |
| 359 | self, input_batch: InputBatch, slot_mappings: torch.Tensor |
| 360 | ) -> None: |
| 361 | """Compile the replay kernels on the buffers they run on: the startup |
| 362 | dummy batches never replay or trim, so they would not launch them.""" |
| 363 | assert self._replay is not None |
| 364 | window, cacheable_groups = self._replay |
| 365 | if not window: |
| 366 | return |
| 367 | _pad_replayed_slots_kernel.warmup( |
| 368 | slot_mappings, |
| 369 | slot_mappings.stride(0), |
| 370 | cacheable_groups, |
| 371 | input_batch.query_start_loc, |
| 372 | input_batch.positions, |
| 373 | self._replay_start, |
| 374 | window, |
| 375 | PAD_SLOT_ID, |
| 376 | NUM_GROUPS=cacheable_groups.numel(), |
| 377 | BLOCK=1024, |
| 378 | grid=(1,), |
| 379 | ) |
| 380 | if self.decoder_replay_layers is None: |
| 381 | return |
| 382 | assert self._kept_slot_mappings is not None |
| 383 | _gather_replay_batch_kernel.warmup( |
| 384 | input_batch.query_start_loc, |
| 385 | self._dropped_before.copy_to_uva(np.zeros(1, dtype=np.int32)), |
| 386 | input_batch.seq_lens, |
| 387 | self._replay_start, |
| 388 | input_batch.positions, |
| 389 | slot_mappings, |
| 390 | slot_mappings.stride(0), |
| 391 | self._kept_rows, |
| 392 | self._kept_query_start_loc, |
| 393 | self._kept_positions, |
| 394 | self._kept_slot_mappings, |
| 395 | self._kept_slot_mappings.stride(0), |
| 396 | self._kept_kv_start, |
| 397 | self.decoder_replay_layers.window, |
| 398 | NUM_GROUPS=slot_mappings.shape[0], |
| 399 | BLOCK=1024, |
| 400 | grid=(1,), |
| 401 | ) |
| 402 | |
| 403 | def _prepare_replay_batch( |
| 404 | self, |
| 405 | input_batch: InputBatch, |
| 406 | cudagraph_mode: CUDAGraphMode, |
| 407 | block_tables: tuple[torch.Tensor, ...], |
| 408 | slot_mappings: torch.Tensor, |
| 409 | attn_groups: list[list[AttentionGroup]], |
| 410 | kv_cache_config: KVCacheConfig, |
| 411 | replay_start: torch.Tensor, |
| 412 | ) -> None: |
| 413 | """Set the replay layers' batch for this forward, or none when nothing |
| 414 | trims. Only eager steps trim: a CUDA graph keeps the layers on the whole |
| 415 | batch, and ranks share the graph mode. Under data parallelism every rank |
| 416 | replays if any does.""" |
| 417 | layers = self.decoder_replay_layers |
| 418 | assert layers is not None |
| 419 | layers.rows = layers.forward_context = None |
| 420 | if cudagraph_mode != CUDAGraphMode.NONE: |
| 421 | return |
| 422 | |
| 423 | num_reqs = input_batch.num_reqs |
| 424 | query_lens = np.diff(input_batch.query_start_loc_np[: num_reqs + 1]) |
| 425 | kept_lens = np.where( |
| 426 | self._keeps_rows_np[input_batch.idx_mapping_np[:num_reqs]], |
| 427 | query_lens, |
| 428 | np.minimum(query_lens, layers.window), |
| 429 | ) |
| 430 | # Dummy batches (captures, an idle DP rank's) are not prefills: they |
| 431 | # keep their rows unless a rank trims. |
| 432 | trims = bool( |
| 433 | (kept_lens < query_lens)[input_batch.is_prefilling_np[:num_reqs]].any() |
| 434 | ) |
| 435 | num_tokens = int(kept_lens.sum()) |
| 436 | dp_metadata = None |
| 437 | if self.vllm_config.parallel_config.data_parallel_size > 1: |
| 438 | trims, dp_metadata = self._agree_across_dp(trims, num_tokens) |
| 439 | if not trims: |
| 440 | return |
| 441 | |
| 442 | kept_batch, kept_slot_mappings = self._kept_input_batch( |
| 443 | input_batch, slot_mappings, replay_start, kept_lens |
| 444 | ) |
| 445 | attn_metadata = super().prepare_attn( |
| 446 | kept_batch, |
| 447 | CUDAGraphMode.NONE, |
| 448 | block_tables, |
| 449 | kept_slot_mappings, |
| 450 | self._replay_groups(attn_groups), |
| 451 | kv_cache_config, |
| 452 | model_specific_attn_metadata=ReplayAttnMetadata( |
| 453 | self._kept_kv_start[:num_reqs] |
| 454 | ), |
| 455 | ) |
| 456 | layers.rows = self._kept_rows[:num_tokens] |
| 457 | layers.forward_context = create_forward_context( |
| 458 | attn_metadata, |
| 459 | self.vllm_config, |
| 460 | dp_metadata=dp_metadata, |
| 461 | slot_mapping=build_slot_mappings_by_layer( |
| 462 | kept_slot_mappings, kv_cache_config |
| 463 | ), |
| 464 | ) |
| 465 | |
| 466 | def _kept_input_batch( |
| 467 | self, |
| 468 | input_batch: InputBatch, |
| 469 | slot_mappings: torch.Tensor, |
| 470 | replay_start: torch.Tensor, |
| 471 | kept_lens: np.ndarray, |
| 472 | ) -> tuple[InputBatch, torch.Tensor]: |
| 473 | """Build the sub-``InputBatch`` of each request's last ``kept_lens`` |
| 474 | rows, and its slot mappings. |
| 475 | |
| 476 | The kept rows follow the device boundaries minus the rows trimmed before |
| 477 | them: adaptive verification resizes the decodes on the GPU alone, and |
| 478 | decodes never trim. Like a microbatch's, the sub-batch describes the |
| 479 | forward only; its sampling fields are carried over and must not be read. |
| 480 | """ |
| 481 | num_reqs = input_batch.num_reqs |
| 482 | num_tokens = int(kept_lens.sum()) |
| 483 | kept_query_start_loc_np = np.zeros(num_reqs + 1, dtype=np.int32) |
| 484 | np.cumsum(kept_lens, out=kept_query_start_loc_np[1:]) |
| 485 | dropped_before = self._dropped_before.copy_to_uva( |
| 486 | input_batch.query_start_loc_np[: num_reqs + 1] - kept_query_start_loc_np |
| 487 | ) |
| 488 | assert self._kept_slot_mappings is not None |
| 489 | kept_slot_mappings = self._kept_slot_mappings[:, :num_tokens] |
| 490 | _gather_replay_batch_kernel[(num_reqs,)]( |
| 491 | input_batch.query_start_loc, |
| 492 | dropped_before, |
| 493 | input_batch.seq_lens, |
| 494 | replay_start, |
| 495 | input_batch.positions, |
| 496 | slot_mappings, |
| 497 | slot_mappings.stride(0), |
| 498 | self._kept_rows, |
| 499 | self._kept_query_start_loc, |
| 500 | self._kept_positions, |
| 501 | kept_slot_mappings, |
| 502 | kept_slot_mappings.stride(0), |
| 503 | self._kept_kv_start, |
| 504 | self.decoder_replay_layers.window, # type: ignore[union-attr] |
| 505 | NUM_GROUPS=slot_mappings.shape[0], |
| 506 | BLOCK=1024, |
| 507 | ) |
| 508 | # Adaptive verification may lengthen the decodes on the device, up to |
| 509 | # the batch's bound. |
| 510 | max_query_len = max(int(kept_lens.max()), input_batch.max_query_len or 0) |
| 511 | kept_batch = replace( |
| 512 | input_batch, |
| 513 | num_tokens=num_tokens, |
| 514 | num_tokens_after_padding=num_tokens, |
| 515 | query_start_loc=self._kept_query_start_loc[: num_reqs + 1], |
| 516 | query_start_loc_np=kept_query_start_loc_np, |
| 517 | max_query_len=max_query_len, |
| 518 | positions=self._kept_positions[:num_tokens], |
| 519 | fast_prefill=None, |
| 520 | ) |
| 521 | return kept_batch, kept_slot_mappings |
| 522 | |
| 523 | def _agree_across_dp( |
| 524 | self, trims: bool, num_tokens: int |
| 525 | ) -> tuple[bool, DPMetadata | None]: |
| 526 | """Whether any rank trims and, then, the replay layers' DP metadata from |
| 527 | every rank's replay token count.""" |
| 528 | parallel_config = self.vllm_config.parallel_config |
| 529 | agreed = torch.zeros(2, parallel_config.data_parallel_size, dtype=torch.int32) |
| 530 | agreed[:, parallel_config.data_parallel_rank] = torch.tensor( |
| 531 | [trims, num_tokens] |
| 532 | ) |
| 533 | if not should_skip_dp_coordination(): |
| 534 | dist.all_reduce(agreed, group=get_dp_group().cpu_group) |
| 535 | if not agreed[0].any(): |
| 536 | return False, None |
| 537 | return True, DPMetadata.make(parallel_config, num_tokens, agreed[1]) |
| 538 | |
| 539 | def _replay_groups( |
| 540 | self, attn_groups: list[list[AttentionGroup]] |
| 541 | ) -> list[list[AttentionGroup]]: |
| 542 | """The attention groups with metadata builders of their own, like each |
| 543 | microbatch's: a builder keeps the metadata it built, and the runner's |
| 544 | hold the batch's.""" |
| 545 | if self._replay_attn_groups is None: |
| 546 | self._replay_attn_groups = [] |
| 547 | for groups in attn_groups: |
| 548 | replay_groups = [] |
| 549 | for group in groups: |
| 550 | replay_group = replace(group, metadata_builders=[]) |
| 551 | replay_group.create_metadata_builders( |
| 552 | self.vllm_config, |
| 553 | self.device, |
| 554 | group.metadata_builders[0].kernel_block_size, |
| 555 | ) |
| 556 | replay_groups.append(replay_group) |
| 557 | self._replay_attn_groups.append(replay_groups) |
| 558 | return self._replay_attn_groups |