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

Signed-off-by: Chen Zhang <[email protected]> Signed-off-by: hzh <[email protected]>
HwwwwwwwH · Jan 22, 2025 · 941a5d5 · 941a5d5
1 parent 055a2b7
commit 941a5d5
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)