| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Decoder-side SWA bounded replay: running the replay layers on their batch. |
| 4 | |
| 5 | The layers past the last KV-source layer own nothing but sliding-window KV, so |
| 6 | in eager prefill steps they run on each request's last ``window`` rows only. |
| 7 | ``DeepseekV41ModelState`` prepares those rows as a sub-batch with attention |
| 8 | metadata and a forward context of its own, like a microbatch; |
| 9 | ``DecoderReplayLayers`` gathers the layer inputs by its rows, runs the layers |
| 10 | under that context and scatters the outputs back to batch rows. Steps that run |
| 11 | in a CUDA graph keep the layers on the whole batch, inside the graph. |
| 12 | """ |
| 13 | |
| 14 | from collections.abc import Callable |
| 15 | |
| 16 | import torch |
| 17 | |
| 18 | from vllm.forward_context import ForwardContext, override_forward_context |
| 19 | from vllm.model_executor.layers.fused_moe.moe_output import MoEOutput |
| 20 | |
| 21 | |
| 22 | class DecoderReplayLayers: |
| 23 | """Runs the replay layers on the step's replay batch. |
| 24 | |
| 25 | ``run_layers`` takes a batch's layer inputs and returns its per-row |
| 26 | outputs. ``row_buffers`` hold per-row results the source layer's indexer |
| 27 | publishes for the layers after it; they are compacted to the replay rows |
| 28 | in place. |
| 29 | """ |
| 30 | |
| 31 | def __init__( |
| 32 | self, |
| 33 | window: int, |
| 34 | run_layers: Callable[..., tuple[torch.Tensor, ...]], |
| 35 | row_buffers: list[torch.Tensor], |
| 36 | ) -> None: |
| 37 | self.window = window |
| 38 | self.run_layers = run_layers |
| 39 | self.row_buffers = row_buffers |
| 40 | # The replay batch, set by the model state every step: its rows of the |
| 41 | # batch and its forward context. None runs the layers on the batch. |
| 42 | self.rows: torch.Tensor | None = None |
| 43 | self.forward_context: ForwardContext | None = None |
| 44 | |
| 45 | def __call__( |
| 46 | self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None |
| 47 | ) -> tuple[torch.Tensor, ...]: |
| 48 | rows = self.rows |
| 49 | if rows is None: |
| 50 | return self.run_layers(hidden_states, *states) |
| 51 | # A trimming step holds a prefill longer than the window, more tokens |
| 52 | # than any step whose MoE leaves its finalize to the next layer. |
| 53 | assert isinstance(hidden_states, torch.Tensor) |
| 54 | num_rows = rows.shape[0] |
| 55 | for buf in self.row_buffers: |
| 56 | buf[:num_rows].copy_(buf.index_select(0, rows)) |
| 57 | with override_forward_context(self.forward_context): |
| 58 | row_outputs = self.run_layers( |
| 59 | hidden_states.index_select(0, rows), |
| 60 | *(None if t is None else t.index_select(0, rows) for t in states), |
| 61 | ) |
| 62 | # The trimmed rows' outputs stay zero; nothing reads them. |
| 63 | num_tokens = hidden_states.shape[0] |
| 64 | return tuple( |
| 65 | out.new_zeros((num_tokens, *out.shape[1:])).index_copy_(0, rows, out) |
| 66 | for out in row_outputs |
| 67 | ) |