Description
The hh_score tensor represents attention scores, where initial values are greater than 0 for valid tokens. The anti_image_mask only penalizes image regions and the recent region, leaving non-image, non-recent (text) tokens unchanged with their original positive scores. After computing top-k positions via torch.topk on the non-recent slice and setting only those to 0 via scatter_, the final mask = self.hh_score >= 0 will mark all text tokens as True (since non top-k text tokens retain >0 scores). This results in retaining at least all text tokens in the KV cache, far exceeding the intended budget of self.hh_size + self.recent_size. Consequently, this causes the KV cache to retain tokens beyond the intended budget, potentially increasing memory usage.
Code Snippet
anti_image_mask = torch.full((self.hh_score.shape[0], self.hh_score.shape[1]), 0)
for i in range(self.image_position.shape[0]):
anti_image_mask[:, image_position[i]:image_position[i]+576] = -65516
anti_image_mask = anti_image_mask.to(device=self.hh_score.device, dtype=self.hh_score.dtype)
anti_image_mask[:, -self.recent_size:] = -65516
self.hh_score = self.hh_score + anti_image_mask
*, keep_topk = torch.topk(self.hh_score[:, :-self.recent_size], self.hh_size, dim=-1)
keep_topk = keep_topk.sort().values
self.hh_score.scatter*(1, keep_topk, 0)
self.hh_score[:, -self.recent_size:] = 0
mask = self.hh_score >= 0
Printed the actual kv cache retained by h2o and lookm respectively, along with the budget. The kv cache retained by lookm exceeds the budget.
Description
The
hh_scoretensor represents attention scores, where initial values are greater than 0 for valid tokens. Theanti_image_maskonly penalizes image regions and the recent region, leaving non-image, non-recent (text) tokens unchanged with their original positive scores. After computing top-k positions viatorch.topkon the non-recent slice and setting only those to 0 viascatter_, the finalmask = self.hh_score >= 0will mark all text tokens as True (since non top-k text tokens retain >0 scores). This results in retaining at least all text tokens in the KV cache, far exceeding the intended budget ofself.hh_size + self.recent_size. Consequently, this causes the KV cache to retain tokens beyond the intended budget, potentially increasing memory usage.Code Snippet
Printed the actual kv cache retained by h2o and lookm respectively, along with the budget. The kv cache retained by lookm exceeds the budget.