Fix CI failure of TypeError: RotaryEmbedding.forward_cpu() got an unexpected keyword argument 'fused_set_kv_buffer_arg' (#11009)

This commit is contained in:
YanbingJiang
2025-09-28 15:31:08 +08:00
committed by GitHub
parent 5519766a4d
commit 66face3598

View File

@@ -206,6 +206,7 @@ class RotaryEmbedding(CustomOp):
query: torch.Tensor,
key: torch.Tensor,
offsets: Optional[torch.Tensor] = None,
fused_set_kv_buffer_arg=None,
) -> Tuple[torch.Tensor, torch.Tensor]:
positions = torch.add(positions, offsets) if offsets is not None else positions
if _is_cpu_amx_available: