[llama3 branch] DraftRetriever Reader panics when there is exactly one matching suffix

[xbfs]

Environment and branch

• Branch: llama3
• Context: When using REST with Llama 3.1 8B and draftretriever for datastore retrieval, the process aborts due to a Rust panic during the search phase.

Problem: Reader::search — panic when only one suffix matches

Symptom:

thread '<unnamed>' panicked at 'called `Option::unwrap()` on a `None` value', src/lib.rs:282:53
Aborted (core dumped)

Cause:
In Reader::search, two binary searches are used to find the start and end of the range of matching suffixes. When the suffix array has exactly one matching suffix for the query, the second binary search may never assign to end_of_indices, so it remains None. The code then calls end_of_indices.unwrap(), which panics.

Fix:
Around line 288, change:

let end_of_indices = end_of_indices.unwrap();

to:

let end_of_indices = end_of_indices.unwrap_or(start_of_indices);

When there is only one match, the end position is then treated as equal to the start position, and we no longer call unwrap() on None.

---

Scope

The change is only in DraftRetriever/src/lib.rs; no other files need to be modified. After applying it, run maturin build --release --strip -i python3.9, reinstall the generated wheel, and verify by loading the datastore and calling Reader::search in the llama3-branch REST workflow.

Summary

Issue	Location	Trigger	Fix
Option::unwrap() on None	Reader::search	Exactly one matching suffix	end_of_indices.unwrap_or(start_of_indices)

It would be helpful to have this fix merged on the llama3 branch so that using Llama 3 + REST does not abort due to this edge case in DraftRetriever.

[zhenyuhe00]

Thank you for your great work on this fix! You are very welcome to open a pull request to merge it.

[xbfs]

Okey, I submitted a PR on #35.

[xbfs]

I also wanted to say I really like the REST work — the retrieval-based speculative decoding idea is clear and the speedup on long sequences is very useful. I was wondering if you have any plans to integrate REST into vllm or sglang down the road. I think a lot of people would benefit from that. Thanks for building and sharing this.😊