From 0fa20ad23eda0dd31398c57e46f13b48683257d2 Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Tue, 1 Apr 2025 17:34:07 -0400 Subject: [PATCH] Fix input triton kernel for eagle --- vllm/v1/spec_decode/eagle.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py index 57c6b652593d..3aaaf34bc79b 100644 --- a/vllm/v1/spec_decode/eagle.py +++ b/vllm/v1/spec_decode/eagle.py @@ -250,13 +250,12 @@ def prepare_input_kernel( num_tokens = end_pos - start_pos 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) for i in tl.range(num_blocks): offset = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) tl.store( out_ptr + start_pos + offset, - indices, + index_start + offset, mask=offset < num_tokens, )