[Misc]: Question: Where can I find getting computed kv cache code on v0.

[hiroki-nishimoto-fixstars]

Hi, I'm trying understanding the work flow of vLLM and I'm intersting in Prefix Caching.

So, I want to know the conditions of prefix caching and who (Scheduler, Executor, Worker, Runner etc. ) get kv cache.

In v1 code, I found it.

vllm/vllm/v1/core/scheduler.py

Lines 215 to 217 in 067fa22

	# Get already-cached tokens.
	computed_blocks, num_computed_tokens = \
	self.kv_cache_manager.get_computed_blocks(request)

But, v0, I failed to find it.

Does any one help me?

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[neurowelt]

From what I understand about the vLLM workflow you should be interested in the model_runner.py file.

I would definitely take a look at the execute_model method:

vllm/vllm/worker/model_runner.py

Lines 1677 to 1684 in 96e06e3

	def execute_model(
	self,
	model_input: ModelInputForGPUWithSamplingMetadata,
	kv_caches: List[torch.Tensor],
	intermediate_tensors: Optional[IntermediateTensors] = None,
	num_steps: int = 1,
	**kwargs,
	) -> Optional[Union[List[SamplerOutput], IntermediateTensors]]:

This is the core method for generating the next token.

The ModelRunner class is used by workers to calculate outputs – a very nice example of that can be seen here:

vllm/vllm/worker/worker_base.py

Lines 420 to 427 in 96e06e3

	output = self.model_runner.execute_model(
	model_input=model_input,
	kv_caches=self.kv_cache[worker_input.virtual_engine]
	if self.kv_cache is not None else None,
	intermediate_tensors=intermediate_tensors,
	num_steps=num_steps,
	**kwargs,
	)

As per prefix caching, in the same file as ModelRunner you can find the ModelInputForGPUBuilder class that implements _compute_for_prefix_cache_hit method:

vllm/vllm/worker/model_runner.py

Lines 537 to 543 in 96e06e3

	def _compute_for_prefix_cache_hit(
	self, inter_data: InterDataForSeqGroup, seq_idx: int,
	seq_group_metadata: SequenceGroupMetadata):
	"""Check if hit prefix cache (i.e., some blocks are already computed).
	If hit, update input tokens and positions to only compute the
	remaining blocks.
	"""

If you follow the methods, you can see exactly where it is used to calculate prefix cache:

1. ModelInputForGPUBuilder is used as _builder_cls in the model runner
   

   vllm/vllm/worker/model_runner.py

   Lines 1103 to 1105 in 96e06e3

   	if hasattr(self, "_builder_cls"):
   	# multi-step model runner does not have `_builder_cls`
   	self.builder = self._builder_cls(weakref.proxy(self))

2. self.builder is used to prepare inputs and add sequence groups for later processing
   

   vllm/vllm/worker/model_runner.py

   Lines 1211 to 1222 in 96e06e3

   	self.builder.prepare(finished_requests_ids)
   	for seq_group_metadata in seq_group_metadata_list:
   	try:
   	self.builder.add_seq_group(seq_group_metadata)
   	except Exception as e:
   	# Raise an exception that tracks the ID of the bad request
   	raise InputProcessingError(seq_group_metadata.request_id,
   	str(e)) from e
   	self.builder.reset_cached_inter_data()
   	return self.builder.build() # type: ignore

3. When adding sequence groups, at the end we call a bunch of per_seq_compute_fns on each sequence
   

   vllm/vllm/worker/model_runner.py

   Lines 756 to 760 in 96e06e3

   	for seq_idx in range(n_seqs):
   	for per_seq_fn in self.per_seq_compute_fns:
   	per_seq_fn(inter_data, seq_idx, seq_group_metadata)
   	for per_seq_group_fn in self.per_seq_group_compute_fns:
   	per_seq_group_fn(inter_data, seq_group_metadata)

4. One of them is the _compute_for_prefix_cache_hit you are looking for
   

   vllm/vllm/worker/model_runner.py

   Lines 438 to 443 in 96e06e3

   	self.per_seq_compute_fns = [
   	self._compute_lens,
   	self._compute_for_prefix_cache_hit,
   	self._compute_for_sliding_window,
   	self._compute_lora_input,
   	]

I hope that clarifies the flow a little bit, please let me know if I missed anything! 😊