[AMD] Fix DSV4 unified attention sink TP slice (#35092)
Co-authored-by: kk <43161300+kkHuang-amd@users.noreply.github.com>
This commit is contained in:
@@ -1699,7 +1699,7 @@ class MQALayer(MqaAttentionBase):
|
|||||||
layer=self.attn_mqa,
|
layer=self.attn_mqa,
|
||||||
forward_batch=forward_batch,
|
forward_batch=forward_batch,
|
||||||
compress_ratio=self.compress_ratio,
|
compress_ratio=self.compress_ratio,
|
||||||
attn_sink=self.attn_sink,
|
attn_sink=attn_sink[: self.n_local_heads],
|
||||||
save_kv_cache=kv is not None,
|
save_kv_cache=kv is not None,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
|
|||||||
Reference in New Issue
Block a user