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

Signed-off-by: Chen Zhang <[email protected]>
vllm-project · Jan 14, 2025 · 1f18adb · 1f18adb
1 parent bb354e6
commit 1f18adb
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)