Skip to content

vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils

Classes:

SpeculatorCudaGraphManager

Bases: CudaGraphManager

CudaGraphManager for draft prefill and decode.

Builds fresh dummy inputs and attention metadata for every warmup and capture pass so that the contents of the shared persistent buffers (e.g. query_start_loc, seq_lens, FA3 scheduler metadata) always match the batch descriptor being captured. Reusing metadata built during an earlier capture would execute kernels with stale buffer contents.

Source code in vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py 
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 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
class SpeculatorCudaGraphManager(CudaGraphManager):
    """CudaGraphManager for draft prefill and decode.

    Builds fresh dummy inputs and attention metadata for every warmup and
    capture pass so that the contents of the shared persistent buffers
    (e.g. query_start_loc, seq_lens, FA3 scheduler metadata) always match
    the batch descriptor being captured. Reusing metadata built during an
    earlier capture would execute kernels with stale buffer contents.
    """

    def specialize_spec_tokens(
        self, desc: BatchExecutionDescriptor, num_speculative_tokens: int
    ) -> BatchExecutionDescriptor:
        if desc.cg_mode != CUDAGraphMode.FULL or desc in self.graphs:
            return desc
        desc = SpeculatorBatchDescriptor(
            **vars(desc), num_speculative_tokens=num_speculative_tokens
        )
        return (
            desc if desc in self.graphs else replace(desc, cg_mode=CUDAGraphMode.NONE)
        )

    def capture(
        self,
        forward_fn: Callable,
        model_state: ModelState,
        input_buffers: InputBuffers,
        block_tables: BlockTables,
        attn_groups: list[list[AttentionGroup]],
        kv_cache_config: KVCacheConfig,
        progress_bar_desc: str = "Capturing CUDA graphs",
        specialize_spec_tokens: bool = False,
    ) -> None:
        if specialize_spec_tokens:
            speculative_config = self.vllm_config.speculative_config
            assert speculative_config is not None
            dense_schedule = build_dynamic_sd_schedule_lookup(
                speculative_config.num_speculative_tokens_per_batch_size,
                vllm_max_batch_size=self.max_num_reqs,
                vllm_num_speculative_tokens=self.vllm_config.num_speculative_tokens,
            )
            self._capture_descs = {
                mode: [
                    SpeculatorBatchDescriptor(**vars(desc), num_speculative_tokens=k)
                    for desc in descs
                    for k in sorted(set(dense_schedule) - {0, 1}, reverse=True)
                ]
                for mode, descs in self._capture_descs.items()
            }

        def create_forward_fn(
            desc: BatchExecutionDescriptor,
            warmup: bool,
        ) -> Callable[[CUDAGraphMode], None]:
            num_tokens = desc.num_tokens
            num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)
            num_tokens_across_dp = (
                torch.full((self.dp_size,), num_tokens, dtype=torch.int32, device="cpu")
                if self.dp_size > 1
                else None
            )
            attn_metadata, slot_mappings = prepare_inputs_to_capture(
                num_reqs,
                num_tokens,
                model_state,
                input_buffers,
                block_tables,
                attn_groups,
                kv_cache_config,
                full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL,
                max_query_len=desc.max_query_len or desc.uniform_token_count,
            )

            forward_kwargs = {}
            if self.decode_query_len == 1:
                forward_kwargs["num_speculative_steps"] = (
                    desc.num_speculative_tokens
                    if isinstance(desc, SpeculatorBatchDescriptor)
                    else self.vllm_config.num_speculative_tokens
                )
            return lambda cg_mode: forward_fn(
                num_reqs,
                num_tokens,
                attn_metadata,
                slot_mappings,
                num_tokens_across_dp,
                cg_mode,
                **forward_kwargs,
            )

        super().capture(create_forward_fn, progress_bar_desc)