Skip to content

[Kernel] mark TorchSDPABackend swap_blocks NotImplementedError #19749

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mgoin merged 1 commit into from
Jun 20, 2025
Merged

[Kernel] mark TorchSDPABackend swap_blocks NotImplementedError #19749

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion vllm/attention/backends/torch_sdpa.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ def swap_blocks(
dst_kv_cache: torch.Tensor,
src_to_dst: torch.Tensor,
) -> None:
PagedAttention.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)
raise NotImplementedError("Swap is not supported in TorchSDPABackend.")
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot Jun 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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.


@staticmethod
def copy_blocks(
Expand Down