[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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

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)