Skip to content

vllm.v1.worker.cpu.model_states.interface

Classes:

ModelSpecificAttnMetadata

Base class for model-specific attention metadata.

Source code in vllm/v1/worker/cpu/model_states/interface.py 
21
22
23
24
25
26
27
28
29
30
31
32
class ModelSpecificAttnMetadata:
    """Base class for model-specific attention metadata."""

    def get_extra_common_attn_kwargs(
        self, kv_cache_group_id: int, num_reqs: int
    ) -> dict[str, Any]:
        return {}

    def get_extra_attn_kwargs(
        self, attn_metadata_builder: Any, num_reqs: int
    ) -> dict[str, Any]:
        return {}

ModelState

Bases: ABC

Methods:

Attributes:

Source code in vllm/v1/worker/cpu/model_states/interface.py 
 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
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
class ModelState(ABC):
    def __init__(
        self,
        vllm_config: VllmConfig,
        model: nn.Module,
        encoder_cache: EncoderCache | None,
        device: torch.device,
    ) -> None:
        self.vllm_config = vllm_config
        self.model_config = vllm_config.model_config
        self.scheduler_config = vllm_config.scheduler_config
        self.model = model
        self.device = device

        self.max_num_reqs = self.scheduler_config.max_num_seqs
        self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
        self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
        self.dtype = self.model_config.dtype

        self.supports_mm_inputs = encoder_cache is not None
        if encoder_cache is not None:
            self.encoder_cache = encoder_cache
            observability_config = vllm_config.observability_config
            self.encoder_runner = EncoderRunner(
                model=self.model,
                max_num_tokens=self.max_num_tokens,
                hidden_size=self.inputs_embeds_size,
                encoder_cache=encoder_cache,
                dtype=self.dtype,
                device=self.device,
                enable_timing=bool(
                    observability_config
                    and observability_config.enable_mm_processor_stats
                ),
            )

    @property
    def max_model_len(self) -> int:
        # Auto-fit can reduce the limit after model state initialization.
        return self.model_config.max_model_len

    def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]:
        from vllm.model_executor.models.interfaces import (
            supports_realtime,
            supports_transcription,
        )
        from vllm.model_executor.models.interfaces_base import is_text_generation_model

        supported_tasks = list[GenerationTask]()
        if is_text_generation_model(self.model):
            supported_tasks.append("generate")
        if supports_transcription(self.model):
            if self.model.supports_transcription_only:
                return ("transcription",)
            supported_tasks.append("transcription")
        if supports_realtime(self.model):
            supported_tasks.append("realtime")
        return tuple(supported_tasks)

    def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
        return None

    def remove_request(self, req_id: str) -> None:
        return None

    def initialize_kv_cache(
        self, kv_cache_config: KVCacheConfig, block_tables: CPUBlockTables
    ) -> None:
        """Hook run after the KV cache tensors are allocated and bound."""
        return None

    def preprocess_state(
        self,
        input_batch: InputBatch,
        block_tables: tuple[torch.Tensor, ...],
        kv_cache_config: KVCacheConfig,
        num_computed_tokens: torch.Tensor,
    ) -> None:
        """Hook run on real batches before the forward pass (after block tables
        are gathered). Used by mamba "align" prefix caching to pre-copy state
        across block boundaries. No-op by default."""
        return None

    @abstractmethod
    def prepare_inputs_embeds(
        self,
        scheduled_encoder_inputs: dict[str, list[int]],
        input_batch: InputBatch,
        req_states: CPURequestState,
    ) -> torch.Tensor | None:
        """Prepare the ``inputs_embeds`` tensor for the current forward pass."""
        raise NotImplementedError

    def dummy_inputs_embeds(self, num_tokens: int) -> torch.Tensor | None:
        """Pre-allocated inputs_embeds buffer for dummy runs (contents unused)."""
        return None

    def execute_mm_encoder(
        self, scheduled_encoder_inputs: dict[str, list[int]]
    ) -> None:
        """Run the multi-modal encoder and cache its outputs by `mm_hash`.

        The encode half of `get_mm_embeddings`, without the gather, for callers
        that run no language model.
        """
        mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs(
            scheduled_encoder_inputs
        )
        if mm_kwargs:
            with self.encoder_runner.timed_encoder_operation(
                scheduled_encoder_inputs.keys()
            ):
                encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
            self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs))

    def gather_mm_embeddings(
        self, input_batch: InputBatch
    ) -> tuple[list[torch.Tensor], torch.Tensor]:
        """Gather cached multimodal embeddings."""
        return self.encoder_runner.gather_mm_embeddings(
            input_batch.req_ids,
            input_batch.num_tokens,
            input_batch.num_scheduled_tokens,
            input_batch.query_start_loc_np,
            input_batch.prefill_len_np,
            input_batch.num_computed_tokens_np,
        )

    @abstractmethod
    def prepare_inputs(
        self, input_batch: InputBatch, req_states: CPURequestState
    ) -> dict[str, Any]:
        raise NotImplementedError

    @abstractmethod
    def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
        raise NotImplementedError

    @abstractmethod
    def prepare_attn(
        self,
        input_batch: InputBatch,
        block_tables: tuple[torch.Tensor, ...],
        slot_mappings: torch.Tensor,
        attn_groups: list[list[AttentionGroup]],
        kv_cache_config: KVCacheConfig,
    ) -> dict[str, Any]:
        raise NotImplementedError

    num_new_sampled_tokens_per_step: int = 1
    """New tokens sampled on each decode step
    (excluding accepted draft tokens, a.k.a num bonus tokens)."""

num_new_sampled_tokens_per_step = 1 class-attribute instance-attribute

New tokens sampled on each decode step (excluding accepted draft tokens, a.k.a num bonus tokens).

dummy_inputs_embeds(num_tokens)

Pre-allocated inputs_embeds buffer for dummy runs (contents unused).

Source code in vllm/v1/worker/cpu/model_states/interface.py 
128
129
130
def dummy_inputs_embeds(self, num_tokens: int) -> torch.Tensor | None:
    """Pre-allocated inputs_embeds buffer for dummy runs (contents unused)."""
    return None

execute_mm_encoder(scheduled_encoder_inputs)

Run the multi-modal encoder and cache its outputs by mm_hash.

The encode half of get_mm_embeddings, without the gather, for callers that run no language model.

Source code in vllm/v1/worker/cpu/model_states/interface.py 
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
def execute_mm_encoder(
    self, scheduled_encoder_inputs: dict[str, list[int]]
) -> None:
    """Run the multi-modal encoder and cache its outputs by `mm_hash`.

    The encode half of `get_mm_embeddings`, without the gather, for callers
    that run no language model.
    """
    mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs(
        scheduled_encoder_inputs
    )
    if mm_kwargs:
        with self.encoder_runner.timed_encoder_operation(
            scheduled_encoder_inputs.keys()
        ):
            encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
        self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs))

gather_mm_embeddings(input_batch)

Gather cached multimodal embeddings.

Source code in vllm/v1/worker/cpu/model_states/interface.py 
150
151
152
153
154
155
156
157
158
159
160
161
def gather_mm_embeddings(
    self, input_batch: InputBatch
) -> tuple[list[torch.Tensor], torch.Tensor]:
    """Gather cached multimodal embeddings."""
    return self.encoder_runner.gather_mm_embeddings(
        input_batch.req_ids,
        input_batch.num_tokens,
        input_batch.num_scheduled_tokens,
        input_batch.query_start_loc_np,
        input_batch.prefill_len_np,
        input_batch.num_computed_tokens_np,
    )

initialize_kv_cache(kv_cache_config, block_tables)

Hook run after the KV cache tensors are allocated and bound.

Source code in vllm/v1/worker/cpu/model_states/interface.py 
100
101
102
103
104
def initialize_kv_cache(
    self, kv_cache_config: KVCacheConfig, block_tables: CPUBlockTables
) -> None:
    """Hook run after the KV cache tensors are allocated and bound."""
    return None

prepare_inputs_embeds(scheduled_encoder_inputs, input_batch, req_states) abstractmethod

Prepare the inputs_embeds tensor for the current forward pass.

Source code in vllm/v1/worker/cpu/model_states/interface.py 
118
119
120
121
122
123
124
125
126
@abstractmethod
def prepare_inputs_embeds(
    self,
    scheduled_encoder_inputs: dict[str, list[int]],
    input_batch: InputBatch,
    req_states: CPURequestState,
) -> torch.Tensor | None:
    """Prepare the ``inputs_embeds`` tensor for the current forward pass."""
    raise NotImplementedError

preprocess_state(input_batch, block_tables, kv_cache_config, num_computed_tokens)

Hook run on real batches before the forward pass (after block tables are gathered). Used by mamba "align" prefix caching to pre-copy state across block boundaries. No-op by default.

Source code in vllm/v1/worker/cpu/model_states/interface.py 
106
107
108
109
110
111
112
113
114
115
116
def preprocess_state(
    self,
    input_batch: InputBatch,
    block_tables: tuple[torch.Tensor, ...],
    kv_cache_config: KVCacheConfig,
    num_computed_tokens: torch.Tensor,
) -> None:
    """Hook run on real batches before the forward pass (after block tables
    are gathered). Used by mamba "align" prefix caching to pre-copy state
    across block boundaries. No-op by default."""
    return None