In prepare_decode():

slot_mappings.append(seq.block_table[-1] * self.block_size + seq.last_block_num_tokens - 1)

I think this is wrong when seq.num_tokens % block_size == 0.

For example, let block_size = 4.

If a sequence already has 8 tokens, then it exactly fills 2 blocks:

• block 10: tokens 1~4
• block 11: tokens 5~8

Now seq.last_block_num_tokens is defined as:

@property
def last_block_num_tokens(self):
    full_blocks = int(math.floor(self.num_tokens / self.block_size))
    return len(self.token_ids[full_blocks * self.block_size : ])

For num_tokens = 8 and block_size = 4:

full_blocks = floor(8 / 4) = 2
len(self.token_ids[2 * 4:]) = len(self.token_ids[8:]) = 0

So seq.last_block_num_tokens = 0.

Then the current formula in prepare_decode() gives:

slot = 11 * 4 + 0 - 1 = 43

But I think this is wrong, because the last token is still in block 11, at the last position, so the slot should be:

11 * 4 + 3 = 47