biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[V1][Spec Decode] Small refactors to improve eagle bookkeeping performance (#18424)

Browse Source 
Signed-off-by: qizixi <qizixi@meta.com>
This commit is contained in:
qizixi
2025-05-23 23:51:22 -07:00
committed by GitHub
parent ec82c3e388
commit d55e446d13
3 changed files with 21 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tests/v1/spec_decode/test_eagle.py
View File

@@ -100,8 +100,12 @@ def test_prepare_inputs():
dtype=torch.int32,
device=device)
# n1 + n2 + n3 - a - b -c
num_tokens = cu_target_query_lens[-1].item() - num_rejected_tokens.sum(
).item()
cu_num_tokens, token_indices = EagleProposer.prepare_inputs(
cu_target_query_lens, num_rejected_tokens)
cu_target_query_lens, num_rejected_tokens, num_tokens)
assert torch.equal(cu_num_tokens, expected_cu_num_tokens)
assert token_indices.shape[0] == expected_cu_num_tokens[-1].item()