Skip to content

vllm.models.deepseek_v41.decoder_replay_layers

Decoder-side SWA bounded replay: running the replay layers on their batch.

Once the last KV-source layer has written every row's KV, it and the layers past it (which own only SWA KV) run on each request's last window rows only. DeepseekV41ModelState prepares those rows as a sub-batch with attention metadata and a forward context of its own, like a microbatch; DecoderReplayLayers gathers the layer inputs by its rows, runs the layers under that context and scatters the outputs back to batch rows. FULL graphs and small PIECEWISE ones keep the layers on the whole batch; larger PIECEWISE graphs break out to the replay batch, which may run in a graph of its own (decoder_replay_cudagraph.py).

Classes:

  • DecoderReplayLayers –

    Runs the replay layers on the step's replay batch.

  • ReplayBatch –

    One forward's replay batch; run_graph is set when it runs in a graph.

DecoderReplayLayers

Runs the replay layers on the step's replay batch.

run_layers takes a batch's layer inputs and returns its per-row outputs; write_batch_kv takes them too and writes the first replay layer's KV for every row. metadata_prefixes are the attention metadata keys the layers read, so the replay batch builds only those. first_swa_prefix keys the first replay layer's SWA metadata, built with the batch's window starts.

Methods:

Source code in vllm/models/deepseek_v41/decoder_replay_layers.py 
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
class DecoderReplayLayers:
    """Runs the replay layers on the step's replay batch.

    ``run_layers`` takes a batch's layer inputs and returns its per-row outputs;
    ``write_batch_kv`` takes them too and writes the first replay layer's KV for
    every row. ``metadata_prefixes`` are the attention metadata keys the layers
    read, so the replay batch builds only those. ``first_swa_prefix`` keys the
    first replay layer's SWA metadata, built with the batch's window starts.
    """

    def __init__(
        self,
        window: int,
        run_layers: Callable[..., tuple[torch.Tensor, ...]],
        write_batch_kv: Callable[..., None],
        metadata_prefixes: set[str],
        first_swa_prefix: str,
    ) -> None:
        self.window = window
        self.run_layers = run_layers
        self.write_batch_kv = write_batch_kv
        self.metadata_prefixes = metadata_prefixes
        self.first_swa_prefix = first_swa_prefix
        # Set by the model state every step; None runs the layers on the batch.
        self.replay_batch: ReplayBatch | None = None
        # Set by the model state with replay graphs: PIECEWISE model graphs of
        # this many tokens or more break out to the replay batch.
        self.trim_threshold: int | None = None
        self._break_outputs: list[torch.Tensor] | None = None

    def __call__(
        self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None
    ) -> tuple[torch.Tensor, ...]:
        if self.uses_replay_batch():
            # Outside the break, so a PIECEWISE graph captures it.
            self.write_batch_kv(hidden_states, *states)
            if in_piecewise_cudagraph():
                return self._run_in_graph_break(hidden_states, *states)
        return self._run(hidden_states, *states)

    def uses_replay_batch(self) -> bool:
        """Whether this forward uses the replay batch: in a PIECEWISE graph by its
        size, fixed at capture (no replay batch is set then); else if one is set."""
        if in_piecewise_cudagraph():
            batch_descriptor = get_forward_context().batch_descriptor
            assert batch_descriptor is not None
            return self.graph_uses_replay_batch(batch_descriptor.num_tokens)
        return self.replay_batch is not None

    def graph_uses_replay_batch(self, num_tokens: int) -> bool:
        """Whether a PIECEWISE graph of ``num_tokens`` breaks out to the replay
        batch; capture and the model state's runtime replay batch share it."""
        return self.trim_threshold is not None and num_tokens >= self.trim_threshold

    @eager_break_during_capture
    def _run_in_graph_break(
        self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None
    ) -> tuple[torch.Tensor, ...]:
        """A graph break: the outputs land in fixed buffers the next segment reads."""
        outputs = self._run(hidden_states, *states)
        if self._break_outputs is None:
            # Sized by the first call: the largest model graph is captured first.
            self._break_outputs = [torch.empty_like(out) for out in outputs]
        return tuple(
            buf[: out.shape[0]].copy_(out)
            for buf, out in zip(self._break_outputs, outputs)
        )

    def _run(
        self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None
    ) -> tuple[torch.Tensor, ...]:
        replay_batch = self.replay_batch
        if replay_batch is None:
            return self.run_layers(hidden_states, *states)
        # Replay batch steps hold more than `window` tokens: no deferred MoE finalize.
        assert isinstance(hidden_states, torch.Tensor)
        rows = replay_batch.rows
        inputs = (hidden_states, *states)
        with override_forward_context(replay_batch.forward_context):
            if replay_batch.run_graph is not None:
                row_outputs = replay_batch.run_graph(rows, inputs)
            else:
                row_outputs = self.run_layers(
                    *(None if t is None else t.index_select(0, rows) for t in inputs)
                )
        # The trimmed rows' outputs stay zero; nothing reads them.
        num_tokens = hidden_states.shape[0]
        return tuple(
            out.new_zeros((num_tokens, *out.shape[1:])).index_copy_(0, rows, out)
            for out in row_outputs
        )

_run_in_graph_break(hidden_states, *states)

A graph break: the outputs land in fixed buffers the next segment reads.

Source code in vllm/models/deepseek_v41/decoder_replay_layers.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
@eager_break_during_capture
def _run_in_graph_break(
    self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None
) -> tuple[torch.Tensor, ...]:
    """A graph break: the outputs land in fixed buffers the next segment reads."""
    outputs = self._run(hidden_states, *states)
    if self._break_outputs is None:
        # Sized by the first call: the largest model graph is captured first.
        self._break_outputs = [torch.empty_like(out) for out in outputs]
    return tuple(
        buf[: out.shape[0]].copy_(out)
        for buf, out in zip(self._break_outputs, outputs)
    )

graph_uses_replay_batch(num_tokens)

Whether a PIECEWISE graph of num_tokens breaks out to the replay batch; capture and the model state's runtime replay batch share it.

Source code in vllm/models/deepseek_v41/decoder_replay_layers.py 
89
90
91
92
def graph_uses_replay_batch(self, num_tokens: int) -> bool:
    """Whether a PIECEWISE graph of ``num_tokens`` breaks out to the replay
    batch; capture and the model state's runtime replay batch share it."""
    return self.trim_threshold is not None and num_tokens >= self.trim_threshold

uses_replay_batch()

Whether this forward uses the replay batch: in a PIECEWISE graph by its size, fixed at capture (no replay batch is set then); else if one is set.

Source code in vllm/models/deepseek_v41/decoder_replay_layers.py 
80
81
82
83
84
85
86
87
def uses_replay_batch(self) -> bool:
    """Whether this forward uses the replay batch: in a PIECEWISE graph by its
    size, fixed at capture (no replay batch is set then); else if one is set."""
    if in_piecewise_cudagraph():
        batch_descriptor = get_forward_context().batch_descriptor
        assert batch_descriptor is not None
        return self.graph_uses_replay_batch(batch_descriptor.num_tokens)
    return self.replay_batch is not None

ReplayBatch dataclass

One forward's replay batch; run_graph is set when it runs in a graph.

Source code in vllm/models/deepseek_v41/decoder_replay_layers.py 
31
32
33
34
35
36
37
@dataclass
class ReplayBatch:
    """One forward's replay batch; ``run_graph`` is set when it runs in a graph."""

    rows: torch.Tensor
    forward_context: ForwardContext
    run_graph: Callable[..., tuple[torch.Tensor, ...]] | None = None