Skip to content

vllm.models.deepseek_v41.nvidia.specialized_decode.runtime.metadata_graphs

CUDA graphs of the attention metadata builds of FULL decode steps.

A FULL CUDA graph replays the forward only: before every replay the runner still builds the attention metadata on the host, which refreshes the persistent buffers the graph reads. For small decode batches that build costs about as much as the replay. AttnMetadataCudaGraphs records, during CUDA graph capture, the build of every uniform decode batch shape that a FULL decode graph serves, and replays it on the steps of that shape, returning the metadata objects captured with it. A step of a shape without a recording builds its metadata eagerly.

Classes:

AttnMetadataCudaGraphs

Records attention metadata builds by graph key during CUDA graph capture, and replays them.

Each key is built once eagerly, warming lazy allocations, kernel compilation and per-stream library state, then recorded with the metadata it returns. The graphs write the builders' buffers at their capture-time addresses: every capture releases and records them again.

Methods:

  • can_record –

    Whether every metadata builder of attn_groups is recordable.

  • capture –

    Record the build of key in the global CUDA graph pool, which the

  • capture_keys –

    The keys of the decode steps the FULL decode graphs serve, largest

  • key –

    The graph key of a step whose metadata build can be recorded, or None.

  • release –

    Drop the recorded builds.

  • replay –

    Replay the build of key and return its metadata, or None if the

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/metadata_graphs.py 
 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
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
class AttnMetadataCudaGraphs:
    """Records attention metadata builds by graph key during CUDA graph
    capture, and replays them.

    Each key is built once eagerly, warming lazy allocations, kernel
    compilation and per-stream library state, then recorded with the metadata
    it returns. The graphs write the builders' buffers at their capture-time
    addresses: every capture releases and records them again.
    """

    def __init__(self, vllm_config: VllmConfig, decode_query_len: int) -> None:
        self._vllm_config = vllm_config
        self._decode_query_len = decode_query_len
        self._graphs: dict[GraphKey, tuple[torch.cuda.CUDAGraph, dict[str, Any]]] = {}
        self._side_streams: list[torch.cuda.Stream] = []
        self._recordable: bool | None = None

    @staticmethod
    def key(input_batch: InputBatch, cudagraph_mode: CUDAGraphMode) -> GraphKey | None:
        """The graph key of a step whose metadata build can be recorded, or None.

        Only FULL steps of uniform decodes qualify; their key fixes every host
        value the build reads. Under adaptive verification (``max_query_len``
        set) the device splits the drafts, so the host query lengths are only
        bounds.
        """
        if (
            cudagraph_mode != CUDAGraphMode.FULL
            or input_batch.has_prefill
            or input_batch.max_query_len is not None
        ):
            return None
        query_lens = input_batch.num_scheduled_tokens[: input_batch.num_reqs]
        if not (query_lens == query_lens[0]).all():
            return None
        return (
            input_batch.num_reqs,
            input_batch.num_reqs_after_padding,
            input_batch.num_tokens,
            input_batch.num_tokens_after_padding,
        )

    def capture_keys(self) -> list[GraphKey]:
        """The keys of the decode steps the FULL decode graphs serve, largest
        first.

        The graphs are the ones ``CudaGraphManager`` captures, and a step runs
        the smallest graph that holds its requests. Under dynamic speculative
        decoding only the widest query length is recorded; steps of the
        narrower tiers build their metadata eagerly.
        """
        config = self._vllm_config.compilation_config
        query_len = self._decode_query_len
        max_reqs = self._vllm_config.get_decode_cudagraph_query_lens(query_len).get(
            query_len, 0
        )
        graph_reqs = [
            n
            for n, _ in decode_cudagraph_capture_shapes(
                config.decode_cudagraph_capture_request_sizes or (),
                {query_len: max_reqs},
                decode_cudagraph_capture_max_tokens(config),
            )
        ]
        keys: list[GraphKey] = []
        num_reqs = 1
        for padded in graph_reqs:
            keys += [
                (n, padded, n * query_len, padded * query_len)
                for n in range(num_reqs, padded + 1)
            ]
            num_reqs = padded + 1
        return keys[::-1]

    def can_record(self, attn_groups: list[list[AttentionGroup]]) -> bool:
        """Whether every metadata builder of ``attn_groups`` is recordable."""
        if self._recordable is None:
            self._recordable = all(
                type(group.get_metadata_builder()) in _RECORDABLE_BUILDERS
                for groups in attn_groups
                for group in groups
            )
        return self._recordable

    def capture(
        self,
        key: GraphKey,
        attn_groups: list[list[AttentionGroup]],
        block_tables: Sequence[torch.Tensor],
        build_fn: BuildFn,
    ) -> dict[str, Any]:
        """Record the build of ``key`` in the global CUDA graph pool, which the
        FULL forward graphs and memory profiling use; returns its metadata."""
        self._forked_build(attn_groups, block_tables, build_fn)
        graph = torch.cuda.CUDAGraph()
        with (
            capture_pool(
                current_platform.get_global_graph_pool(), self._vllm_config
            ) as pool,
            torch.cuda.graph(graph, pool),
        ):
            attn_metadata = self._forked_build(attn_groups, block_tables, build_fn)
        self._graphs[key] = (graph, attn_metadata)
        return attn_metadata

    def replay(self, key: GraphKey) -> dict[str, Any] | None:
        """Replay the build of ``key`` and return its metadata, or None if the
        key was not recorded."""
        entry = self._graphs.get(key)
        if entry is None:
            logger.warning_once(
                "DeepSeek-V4.1 specialized decode: a FULL decode step has no "
                "recorded attention-metadata build; building it eagerly."
            )
            return None
        graph, attn_metadata = entry
        graph.replay()
        return attn_metadata

    def _forked_build(
        self,
        attn_groups: list[list[AttentionGroup]],
        block_tables: Sequence[torch.Tensor],
        build_fn: BuildFn,
    ) -> dict[str, Any]:
        """Build the attention groups as concurrent branches on side streams.

        The groups share only read-only inputs, with two exceptions. The sparse
        MLA and compressor builders clamp their block table in place, so the
        groups of one block table storage build on one branch, in serial order;
        the groups that only read the table (``_BLOCK_TABLE_READERS``) then
        build on branches of their own, each started after that branch. And
        the first group that needs the token-to-request map builds it into its
        own buffer, which later groups reuse. Each branch builds its own map;
        the serial build's producer is also the first of its branch to need
        one, so it fills the buffer that the FULL forward graphs read.
        """
        tables: dict[int, tuple[Branch, Branch]] = {}
        for i, groups in enumerate(attn_groups):
            if groups:
                storage = block_tables[i].untyped_storage().data_ptr()
                writers, readers = tables.setdefault(storage, ([], []))
                for group in groups:
                    reads = type(group.get_metadata_builder()) in _BLOCK_TABLE_READERS
                    (readers if reads else writers).append((i, group))
        # Each branch with the index of the branch it starts after, if any.
        branches: list[tuple[Branch, int | None]] = []
        for writers, readers in tables.values():
            after = None
            if writers:
                after = len(branches)
                branches.append((writers, None))
            branches += [([reader], after) for reader in readers]
        if len(branches) < 2:
            return build_fn(attn_groups)
        # Reused, so a warm-up creates per-stream state (e.g. DeepGEMM's) on the
        # streams its capture records.
        while len(self._side_streams) < len(branches):
            self._side_streams.append(torch.cuda.Stream())
        main_stream = current_stream()
        built: dict[str, Any] = {}
        for stream, (branch, after) in zip(self._side_streams, branches):
            stream.wait_stream(
                main_stream if after is None else self._side_streams[after]
            )
            with torch.cuda.stream(stream):
                built.update(
                    build_fn(
                        [
                            [group for j, group in branch if j == i]
                            for i in range(len(attn_groups))
                        ]
                    )
                )
        for stream in self._side_streams[: len(branches)]:
            main_stream.wait_stream(stream)
        # The serial build's layer order.
        return {
            name: built[name]
            for groups in attn_groups
            for group in groups
            for name in group.layer_names
        }

    def release(self) -> None:
        """Drop the recorded builds."""
        self._graphs.clear()
        self._recordable = None

_forked_build(attn_groups, block_tables, build_fn)

Build the attention groups as concurrent branches on side streams.

The groups share only read-only inputs, with two exceptions. The sparse MLA and compressor builders clamp their block table in place, so the groups of one block table storage build on one branch, in serial order; the groups that only read the table (_BLOCK_TABLE_READERS) then build on branches of their own, each started after that branch. And the first group that needs the token-to-request map builds it into its own buffer, which later groups reuse. Each branch builds its own map; the serial build's producer is also the first of its branch to need one, so it fills the buffer that the FULL forward graphs read.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/metadata_graphs.py 
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
def _forked_build(
    self,
    attn_groups: list[list[AttentionGroup]],
    block_tables: Sequence[torch.Tensor],
    build_fn: BuildFn,
) -> dict[str, Any]:
    """Build the attention groups as concurrent branches on side streams.

    The groups share only read-only inputs, with two exceptions. The sparse
    MLA and compressor builders clamp their block table in place, so the
    groups of one block table storage build on one branch, in serial order;
    the groups that only read the table (``_BLOCK_TABLE_READERS``) then
    build on branches of their own, each started after that branch. And
    the first group that needs the token-to-request map builds it into its
    own buffer, which later groups reuse. Each branch builds its own map;
    the serial build's producer is also the first of its branch to need
    one, so it fills the buffer that the FULL forward graphs read.
    """
    tables: dict[int, tuple[Branch, Branch]] = {}
    for i, groups in enumerate(attn_groups):
        if groups:
            storage = block_tables[i].untyped_storage().data_ptr()
            writers, readers = tables.setdefault(storage, ([], []))
            for group in groups:
                reads = type(group.get_metadata_builder()) in _BLOCK_TABLE_READERS
                (readers if reads else writers).append((i, group))
    # Each branch with the index of the branch it starts after, if any.
    branches: list[tuple[Branch, int | None]] = []
    for writers, readers in tables.values():
        after = None
        if writers:
            after = len(branches)
            branches.append((writers, None))
        branches += [([reader], after) for reader in readers]
    if len(branches) < 2:
        return build_fn(attn_groups)
    # Reused, so a warm-up creates per-stream state (e.g. DeepGEMM's) on the
    # streams its capture records.
    while len(self._side_streams) < len(branches):
        self._side_streams.append(torch.cuda.Stream())
    main_stream = current_stream()
    built: dict[str, Any] = {}
    for stream, (branch, after) in zip(self._side_streams, branches):
        stream.wait_stream(
            main_stream if after is None else self._side_streams[after]
        )
        with torch.cuda.stream(stream):
            built.update(
                build_fn(
                    [
                        [group for j, group in branch if j == i]
                        for i in range(len(attn_groups))
                    ]
                )
            )
    for stream in self._side_streams[: len(branches)]:
        main_stream.wait_stream(stream)
    # The serial build's layer order.
    return {
        name: built[name]
        for groups in attn_groups
        for group in groups
        for name in group.layer_names
    }

can_record(attn_groups)

Whether every metadata builder of attn_groups is recordable.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/metadata_graphs.py 
139
140
141
142
143
144
145
146
147
def can_record(self, attn_groups: list[list[AttentionGroup]]) -> bool:
    """Whether every metadata builder of ``attn_groups`` is recordable."""
    if self._recordable is None:
        self._recordable = all(
            type(group.get_metadata_builder()) in _RECORDABLE_BUILDERS
            for groups in attn_groups
            for group in groups
        )
    return self._recordable

capture(key, attn_groups, block_tables, build_fn)

Record the build of key in the global CUDA graph pool, which the FULL forward graphs and memory profiling use; returns its metadata.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/metadata_graphs.py 
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
def capture(
    self,
    key: GraphKey,
    attn_groups: list[list[AttentionGroup]],
    block_tables: Sequence[torch.Tensor],
    build_fn: BuildFn,
) -> dict[str, Any]:
    """Record the build of ``key`` in the global CUDA graph pool, which the
    FULL forward graphs and memory profiling use; returns its metadata."""
    self._forked_build(attn_groups, block_tables, build_fn)
    graph = torch.cuda.CUDAGraph()
    with (
        capture_pool(
            current_platform.get_global_graph_pool(), self._vllm_config
        ) as pool,
        torch.cuda.graph(graph, pool),
    ):
        attn_metadata = self._forked_build(attn_groups, block_tables, build_fn)
    self._graphs[key] = (graph, attn_metadata)
    return attn_metadata

capture_keys()

The keys of the decode steps the FULL decode graphs serve, largest first.

The graphs are the ones CudaGraphManager captures, and a step runs the smallest graph that holds its requests. Under dynamic speculative decoding only the widest query length is recorded; steps of the narrower tiers build their metadata eagerly.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/metadata_graphs.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def capture_keys(self) -> list[GraphKey]:
    """The keys of the decode steps the FULL decode graphs serve, largest
    first.

    The graphs are the ones ``CudaGraphManager`` captures, and a step runs
    the smallest graph that holds its requests. Under dynamic speculative
    decoding only the widest query length is recorded; steps of the
    narrower tiers build their metadata eagerly.
    """
    config = self._vllm_config.compilation_config
    query_len = self._decode_query_len
    max_reqs = self._vllm_config.get_decode_cudagraph_query_lens(query_len).get(
        query_len, 0
    )
    graph_reqs = [
        n
        for n, _ in decode_cudagraph_capture_shapes(
            config.decode_cudagraph_capture_request_sizes or (),
            {query_len: max_reqs},
            decode_cudagraph_capture_max_tokens(config),
        )
    ]
    keys: list[GraphKey] = []
    num_reqs = 1
    for padded in graph_reqs:
        keys += [
            (n, padded, n * query_len, padded * query_len)
            for n in range(num_reqs, padded + 1)
        ]
        num_reqs = padded + 1
    return keys[::-1]

key(input_batch, cudagraph_mode) staticmethod

The graph key of a step whose metadata build can be recorded, or None.

Only FULL steps of uniform decodes qualify; their key fixes every host value the build reads. Under adaptive verification (max_query_len set) the device splits the drafts, so the host query lengths are only bounds.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/metadata_graphs.py 
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
@staticmethod
def key(input_batch: InputBatch, cudagraph_mode: CUDAGraphMode) -> GraphKey | None:
    """The graph key of a step whose metadata build can be recorded, or None.

    Only FULL steps of uniform decodes qualify; their key fixes every host
    value the build reads. Under adaptive verification (``max_query_len``
    set) the device splits the drafts, so the host query lengths are only
    bounds.
    """
    if (
        cudagraph_mode != CUDAGraphMode.FULL
        or input_batch.has_prefill
        or input_batch.max_query_len is not None
    ):
        return None
    query_lens = input_batch.num_scheduled_tokens[: input_batch.num_reqs]
    if not (query_lens == query_lens[0]).all():
        return None
    return (
        input_batch.num_reqs,
        input_batch.num_reqs_after_padding,
        input_batch.num_tokens,
        input_batch.num_tokens_after_padding,
    )

release()

Drop the recorded builds.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/metadata_graphs.py 
249
250
251
252
def release(self) -> None:
    """Drop the recorded builds."""
    self._graphs.clear()
    self._recordable = None

replay(key)

Replay the build of key and return its metadata, or None if the key was not recorded.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/metadata_graphs.py 
170
171
172
173
174
175
176
177
178
179
180
181
182
def replay(self, key: GraphKey) -> dict[str, Any] | None:
    """Replay the build of ``key`` and return its metadata, or None if the
    key was not recorded."""
    entry = self._graphs.get(key)
    if entry is None:
        logger.warning_once(
            "DeepSeek-V4.1 specialized decode: a FULL decode step has no "
            "recorded attention-metadata build; building it eagerly."
        )
        return None
    graph, attn_metadata = entry
    graph.replay()
    return attn_metadata