[v1] Remove bind_kv_cache and self.kv_cache in model runner

[heheda12345]

As the kv_caches are not needed by model, we can remove it from model runner now and remove the complex bind_kv_cache function.

However, as self.kv_caches is still used by tpu backend, I move the bind_kv_cache into tpu_model_runner.

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

🚀

[heheda12345]

CC @robertgshaw2-redhat Can you take a look at the TPU changes?

[comaniac]

This file shouldn't be here (in vllm/)?

[comaniac]

This is still a kind of "binding"? To unify the interface for runners, would the following be better?

bind_kv_cache(
    kv_caches,
    self.vllm_config.compilation_config.static_forward_context,
)

And in TPU model runner:

bind_kv_cache(
    kv_caches,
    self.vllm_config.compilation_config.static_forward_context,
    self.kv_caches,
)

So we could have

def bind_kv_cache(
    kv_caches: dict[str, torch.Tensor],
    forward_context: dict[str, "Attention"],
    runner_kv_caches: Optional[list[torch.Tensor]] = None,
) -> None:
    # Bind runner kv caches. Also comment
    # this is only used by TPU for now.
    if runner_kv_caches is not None:
        ...

    # Bind forward_context.
    ...

If TPU is the only model runner (now and future) that needs this, alternatively we could have a separate utility:

def bind_runner_kv_caches(
    kv_caches: dict[str, torch.Tensor],
    runner_kv_caches: list[torch.Tensor],
)
    ...

And in TPU model runner:

bind_kv_cache(
    kv_caches,
    self.vllm_config.compilation_config.static_forward_context,
)
bind_runner_kv_caches(
    kv_caches,
    self.kv_caches,
)

[heheda12345]

In tpu_model_runner, self.kv_caches is only used for detecting whether it is profile run and things like num_kv_heads, num_blocks, block_size. I think they can be removed in the future if we do some refactor on ModelWrapperV1. But I'm not very familiar with the tpu backend.

The binding to forward context is quite simple, so I prefer to implement it in initialize_kv_cache instead of calling a new function if we can remove the binding to self.kv_caches.

[comaniac]

Ok if this is confirmed please add TODO/FIXME to bind_kv_cache in TPU runner and we can keep as it is.

[vanbasten23]

Actually @heheda12345 , could you check if v1's tpu_model_runner is really using the self.kv_caches? I did a quick check. It seems that

vllm/vllm/v1/worker/tpu_model_runner.py

Line 823 in 04421df

kv_caches: list[tuple[torch.Tensor, torch.Tensor]],

never uses it, even if it's listed as an argument.

[heheda12345]

It is used in class ModelWrapperV1(nn.Module):, and we are trying to remove it in #14309 before this PR.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @heheda12345.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork