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

Signed-off-by: Chen Zhang <zhangch99@outlook.com>
This commit is contained in:
Chen Zhang
2025-01-14 20:59:32 +08:00
committed by GitHub
parent bb354e6b2d
commit 1f18adb245

View File

@@ -134,8 +134,8 @@ class Attention(nn.Module):
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)