Override attention metadata for fast prefill in some KV sharing setups

[sarckk]

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

To support YOCO prefill skip for cross-decoder layers in #19719, we need to propagate logits_indices that has been padded to align with cudagraph captured batch sizes as well as size of the original logits_indices.

To do this, we can identify which layers are eligible, store these layer names in truncated_prefill_eligible_layers and dynamically build a subclass of attention metadata for these layers that inherit from the original attn metadata class but adds two new fields: logits_indices_padded and num_logits_indices.

Test Plan

unit test:

pytest tests/v1/e2e/test_kv_sharing_fast_prefill.py::test_kv_sharing_fast_prefill

run with new flag:

vllm serve google/gemma-3n-E2B-it --disable-log-requests --kv-sharing-fast-prefill

Test Result

Unit tests pass

See new warning msg:

WARNING 07-29 14:06:05 [config.py:1739] --kv-sharing-fast-prefill is currently work in progress and not functional yet (i.e. no prefill savings)

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

🚀

[gemini-code-assist]

Code Review

This pull request introduces an option to compute and propagate padded logits_indices to the model's forward pass, controlled by the VLLM_COMPUTE_PADDED_LOGITS_INDICES environment variable. The changes span environment variable handling, forward context management, and the GPU model runner logic. My review found a critical typo in the environment variable name that would prevent the feature from working, and a potential IndexError when handling an empty logits_indices tensor. Both issues include code suggestions for fixes.

[heheda12345]

I'm thinking of adding the final flag enable_kv_sharing_truncated_prefill to this PR instead of introducing a temporary environment variable. We can add a warning to show that this feature is WIP.

I think a better granularity to split PRs is to compute a real gen_indices_padded in this PR.
Then you can test it with Gemma-3n which has YOCO. Can you add a simple test based on this model and update it gradually during implementing more small PRs? For this PR, I think we can run the model as normal, and returns a tensor with

second_hidden_states =  hidden_states[gen_indices_padded],
second_residual = residual[gen_indices_padded],
hidden_states = torch.randn()
residual = torch.randn()
hidden_states[gen_indices] = second_hidden_states[:num_gen_indices]
residual[gen_indices] = second_residual[:num_gen_indices]

[heheda12345]

Thanks for your update. Do we need to pass the args via LLM class? Other comments are some small nits.

[sarckk]

Do we need to pass the args via LLM class?

You're right, we don't need it. Updated to address comments.

[heheda12345]

LGTM! Thank you very much.

[DarkLight1337]

test_kv_sharing_fast_prefill is failing on main, can you look into it?

[sarckk]

test_kv_sharing_fast_prefill is failing on main, can you look into it?

taking a look now (EDIT: seems test is flaky, partial fix in #22038)