[Kernel] Revert the API change of Attention.forward (vllm-project#12038)

Signed-off-by: Chen Zhang <zhangch99@outlook.com> Signed-off-by: ice-tong <xych6@outlook.com>
ice-tong · Jan 18, 2025 · 9e9ca03 · 9e9ca03
1 parent f9a5b08
commit 9e9ca03
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/vllm/attention/layer.py b/vllm/attention/layer.py
@@ -134,8 +134,8 @@ def forward(
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        _kv_cache: torch.Tensor,
-        _attn_metadata: AttentionMetadata,
+        kv_cache: torch.Tensor,
+        attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
         if self.use_output:
             output = torch.empty_like(query)