[Kernel] mark TorchSDPABackend swap_blocks NotImplementedError

[andyxning]

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

TorchSDPABackend is only used with cpu device and for now and cpu does not support cache swap operations as stated in:

vllm/vllm/worker/cpu_worker.py

Lines 91 to 95 in 5a1c2e1

	def swap_in(self, src_to_dst: Dict[int, int]) -> None:
	raise NotImplementedError("Swap is not supported in CPUCacheEngine.")
	def swap_out(self, src_to_dst: Dict[int, int]) -> None:
	raise NotImplementedError("Swap is not supported in CPUCacheEngine.")

So, TorchSDPABackend can raise NotImplementedError.

Test Plan

Test Result

(Optional) Documentation Update

[gemini-code-assist]

Summary of Changes

Hello @andyxning, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request clarifies the behavior of the TorchSDPABackend by explicitly raising a NotImplementedError for the swap_blocks method. This is a necessary change because the TorchSDPABackend is currently designed for CPU usage, and CPU environments do not support the cache swap operations that this method implies. The change prevents unexpected behavior and clearly communicates the current limitations of this specific backend.

Highlights

• Error Handling: Explicitly marked the swap_blocks method within the TorchSDPABackend as NotImplementedError.
• Backend Limitation: This change reflects that the TorchSDPABackend is currently CPU-only, and CPU devices do not support the cache swap operations that this method would typically handle.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[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 correctly addresses an issue with the TorchSDPABackend by explicitly marking its swap_blocks method with a NotImplementedError. This change aligns with the documented limitations of CPU workers, which do not support swap operations. The code modification itself is clear and appropriate.

The main feedback provided focuses on improving the PR process by suggesting the inclusion of a test plan and a corresponding unit test to verify the NotImplementedError is raised as expected. Additionally, a reminder to complete the PR description checklist was given. These process improvements aim to enhance the robustness and maintainability of the codebase.

[gemini-code-assist]

This change to raise NotImplementedError is a good improvement, accurately reflecting that swap operations are not supported for TorchSDPABackend on CPU. This aligns with the rationale provided in the PR description and the linked comments in vllm/worker/cpu_worker.py.

To further enhance the robustness of this change and adherence to PR best practices:

1. Test Plan: The PR description's 'Test Plan' section is currently empty. It would be beneficial to include a test plan. For this specific change, a unit test could be added to verify that calling TorchSDPABackend.swap_blocks() indeed raises NotImplementedError with the expected message. This ensures the intended behavior is explicitly tested and maintained through future refactorings.
   Example test sketch:

   import pytest
   import torch
   from vllm.attention.backends.torch_sdpa import TorchSDPABackend
   
   def test_torch_sdpa_swap_blocks_not_implemented():
       # It's a static method, so no instance is strictly needed for calling,
       # but let's assume it might be called via an instance context in some cases.
       # If it's always called statically, TorchSDPABackend.swap_blocks(...) is fine.
       src_kv_cache = torch.empty(0)
       dst_kv_cache = torch.empty(0)
       src_to_dst = torch.empty(0)
       with pytest.raises(NotImplementedError, match="Swap is not supported in TorchSDPABackend."):
           TorchSDPABackend.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)

2. PR Description Checklist: Please also ensure the checklist items at the top of the PR description template are reviewed and ticked as appropriate.

Adding a unit test would provide greater confidence in this change and help complete the PR description more thoroughly.