Support DraftRetriever datastore read/write for large vocab sizes (i.e. llama3+) and REST inference for llama3

[scandukuri]

Here we make the necessary changes to read and write suffixes to memory and file for large tokenizers; the original implementation only supported token IDs up to Rust u16::MAX (65,535). Crucially, using Rust i32 for reading and writing individual token IDs (instead of u16 originally) allows the tool to support token IDs of up to Rust i32::MAX (2,147,483,647), and still allows negative placeholder IDs for padding in the implementation like -2.

We also make the necessary adjustments torest/modeling_llama_kv.py to support llama3 inference with REST, adapted from the transformers implementation. Importantly, we need to repeat_kv for new keys and values before concatenating them to cached keys and values, in order to avoid shape mismatches (by only adding a single query group to all cached keys and values).

Including the old PR here for reference. Thanks! 🦙

[scandukuri]

@zhenyuhe00 Here's a new PR with a separate branch and added changes to modeling_llama_kv.py - it seems like the base branch may still be main, though. Do you know how I can PR my fork's llama3 changes to create a new llama3 branch in the original codebase?

[zhenyuhe00]

Hi, I appreciate your effort! Since the repository only has the "main" branch, I added a new branch called "llama3" just now. It could be merged into the "llama3" branch.