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 
 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
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
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.

        As in ``CudaGraphManager``, each capture size rounds up to whole
        requests of the decode query length, and a step runs the smallest graph
        that holds its requests.
        """
        config = self._vllm_config.compilation_config
        query_len = self._decode_query_len
        max_tokens = min(
            self._vllm_config.scheduler_config.max_num_seqs * query_len,
            config.max_cudagraph_capture_size,
        )
        graph_reqs = sorted(
            {
                round_up(size, query_len) // query_len
                for size in config.cudagraph_capture_sizes or ()
                if round_up(size, query_len) <= max_tokens
            }
        )
        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 KV-cache groups as concurrent branches on side streams.

        The groups share only read-only inputs, with two exceptions. The sparse
        MLA builder clamps its block table in place, so groups whose block
        tables share storage stay on one branch, in serial order. 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.
        """
        branches: dict[int, list[int]] = {}
        for i, groups in enumerate(attn_groups):
            if groups:
                storage = block_tables[i].untyped_storage().data_ptr()
                branches.setdefault(storage, []).append(i)
        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 in zip(self._side_streams, branches.values()):
            stream.wait_stream(main_stream)
            with torch.cuda.stream(stream):
                built.update(
                    build_fn(
                        [g if i in branch else [] for i, g in enumerate(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 KV-cache groups as concurrent branches on side streams.

The groups share only read-only inputs, with two exceptions. The sparse MLA builder clamps its block table in place, so groups whose block tables share storage stay on one branch, in serial order. 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 
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
def _forked_build(
    self,
    attn_groups: list[list[AttentionGroup]],
    block_tables: Sequence[torch.Tensor],
    build_fn: BuildFn,
) -> dict[str, Any]:
    """Build the KV-cache groups as concurrent branches on side streams.

    The groups share only read-only inputs, with two exceptions. The sparse
    MLA builder clamps its block table in place, so groups whose block
    tables share storage stay on one branch, in serial order. 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.
    """
    branches: dict[int, list[int]] = {}
    for i, groups in enumerate(attn_groups):
        if groups:
            storage = block_tables[i].untyped_storage().data_ptr()
            branches.setdefault(storage, []).append(i)
    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 in zip(self._side_streams, branches.values()):
        stream.wait_stream(main_stream)
        with torch.cuda.stream(stream):
            built.update(
                build_fn(
                    [g if i in branch else [] for i, g in enumerate(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 
129
130
131
132
133
134
135
136
137
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 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
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.

As in CudaGraphManager, each capture size rounds up to whole requests of the decode query length, and a step runs the smallest graph that holds its requests.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/metadata_graphs.py 
 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
def capture_keys(self) -> list[GraphKey]:
    """The keys of the decode steps the FULL decode graphs serve, largest
    first.

    As in ``CudaGraphManager``, each capture size rounds up to whole
    requests of the decode query length, and a step runs the smallest graph
    that holds its requests.
    """
    config = self._vllm_config.compilation_config
    query_len = self._decode_query_len
    max_tokens = min(
        self._vllm_config.scheduler_config.max_num_seqs * query_len,
        config.max_cudagraph_capture_size,
    )
    graph_reqs = sorted(
        {
            round_up(size, query_len) // query_len
            for size in config.cudagraph_capture_sizes or ()
            if round_up(size, query_len) <= max_tokens
        }
    )
    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 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
@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 
221
222
223
224
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 
160
161
162
163
164
165
166
167
168
169
170
171
172
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