misc: correct the int data type for token ids and indices

[xiezhq-hermann]

Motivation

The data type for token ids and KV cache indices are supposed to be int32 consistently instead of int64.

Modification

Simply corrected the misuse of torch.int64.

Checklist

1. Ensure pre-commit pre-commit run --all-files or other linting tools are used to fix potential lint issues.
2. Confirm that modifications are covered by complete unit tests. If not, please add more unit tests for correctness.
3. Modify documentation as needed, such as docstrings or example tutorials.