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

[Spec Decode] Fix input triton kernel for eagle (#15909)

Browse Source
This commit is contained in:
Ekagra Ranjan 2025-04-01 21:15:14 -04:00 committed by GitHub
parent 58f5a59769
commit 24b7fb455a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
vllm/v1/spec_decode/eagle.py
View File

@ -250,13 +250,12 @@ def prepare_input_kernel(
num_tokens = end_pos - start_pos num_tokens = end_pos - start_pos
index_start = tl.load(cu_query_lens_ptr + pid) index_start = tl.load(cu_query_lens_ptr + pid)
indices = index_start + tl.arange(0, BLOCK_SIZE)
num_blocks = tl.cdiv(num_tokens, BLOCK_SIZE) num_blocks = tl.cdiv(num_tokens, BLOCK_SIZE)
for i in tl.range(num_blocks): for i in tl.range(num_blocks):
offset = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) offset = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
tl.store( tl.store(
out_ptr + start_pos + offset, out_ptr + start_pos + offset,
indices, index_start + offset,
mask=offset < num_tokens, mask=offset < num_tokens,
) )