[V1] TPU - Remove self.kv_caches

[alexm-redhat]

This PR removes self.kv_caches from the tpu_model_runner.py in V1, so that @heheda12345 #14098 can cleanly land.

@mgoin @NickLucche feel free to make a pass.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[mgoin]

Please revert this - we should make a examples/offline_inference/tpu/ folder to keep this

[heheda12345]

If we add is_profile_run to forward_context, we need to change other backends to pass the is_profile_run arg. Can we achieve it by passing the attributes to ModelWrapperV1 like the pseudo code?

[heheda12345]

Is it possible to implement ModelWrapperV1 like this?

class ModelWrapperV1(nn.Module):

     def __init__(self, model: nn.Module, num_kv_heads, num_blocks, block_size):
         super().__init__()
         self.model = model
         self.num_kv_heads = num_kv_heads
         ...

     def forward(
         self,
         input_ids: torch.Tensor,
         positions: torch.Tensor,
         inputs_embeds: Optional[torch.Tensor] = None,
         is_profile_run: bool,
     ) -> torch.Tensor:
         if not is_profile_run:
              num_kv_heads = self.num_kv_heads
              ...

class TPUModelRunner:
    def _dummy_run(
           self,
           num_tokens: int,
           is_profile_run: bool,
     ) -> None:
         self.model.forward(..., is_profile_run=is_profile_run)

[alexm-redhat]

@heheda12345 this is not possible because num_blocks is not known until determine_num_available_blocks is done and initialize_kv_cache is executed.

[alexm-redhat]

@heheda12345 the is_profile_run arg is set to False by default, so it should not be necessary to pass this parameter explicitly to set_forward_context(..) function in other backends. Is there any specific code example where you would need to specify it explicitly?