maybe a bug in `update_kv` function

[HarryWu99]

SnapKV/snapkv_utils.py

Line 50 in ea655b1

attention_mask = mask[None, None, :, :]

In update_kv function, instead of using the function's arguments attention_mask, this variable is overridden.

[leeyeehoo]

Just refactored the codebase. I think the arg is dummy and will remove it in the future update. You can test to see if there is any bug! Thx!