[v1] Pass BlockTable and KVCacheSpec to AttentionMetadataBuilders

[heheda12345]

Should merge after #17394

Hybrid allocator will need to build attention metadata for each kv cache group because different kv cache groups may have different attention type and block_table. To achieve that, we will introduce one AttentionMetadataBuilder and one BlockTable for each group.

To prepare for this, this PR makes AttentionMetadataBuilder to access its block_table and KVCacheSpec, instead of reading from model_runner.

And as slot_mapping will also be different for different kv cache groups, this pr moves the slot_mapping_cpu tensor from runner to BlockTable.

Splitted from #16101

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

🚀

[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

[WoosukKwon]

LGTM!

[WoosukKwon]

@heheda12345 Please fix the CI failure 😅

[tjtanaa]

@heheda12345 may I know if the pre-commit check caught this?

the rename of the block_table to block_table_tensor in _build_decode function had broken AiterMLAMetadataBuilder._build_decode

If it doesn't then there might need some investigation to find out the reason why it is not caught in the CI.

[heheda12345]

@tjtanaa An interesting question. I think pre-commit can pass because we have MLACommonMetadataBuilder._build_decode with the correct signature. Maybe we need more strict type annotations.
A small reproduce script:

class A:

    def func(self, a: int, b: int) -> None:
        pass

    def func2(self, a: int, b: int) -> None:
        self.func(a, b=b)


class B(A):

    def func(self, a: int, bb: int) -> None:
        pass


# B().func(1, b=2) # type error
B().func2(1, 2)  # pass type check but crash like you mentioned.

[chenyang78]

Looks like this commit caused regression to the FlashInfer backend (with the FlashInfer's latest commit: 25fb40) at least on GB200 and B200. With this commit, the following command failed with CUDA out-of-memory failures:

$ VLLM_USE_V1=1 VLLM_ATTENTION_BACKEND=FLASHINFER lm_eval --model vllm --model_args pretrained=meta-llama/Llama-3.1-8B-Instruct --trust_remote_code --tasks gsm8k --num_fewshot 5 --batch_size auto
...
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 4.50 GiB. ..

It worked fine right before this commit.

[chenyang78]

The same CUDA out-of-memory failure also occurred on H100 with the same command above, i.e.

$ VLLM_USE_V1=1 VLLM_ATTENTION_BACKEND=FLASHINFER lm_eval --model vllm --model_args pretrained=meta-llama/Llama-3.1-8B-Instruct --trust_remote_code --tasks gsm8k --num_fewshot 5 --batch_size auto 
...
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.12 GiB. GPU 0 has a total capacity of 95.00 GiB of which 1.01 GiB is free.

BTW, I was using CUDA 12.8. Thanks!