A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 16:00 UTC
[Spec Decode] Fix input triton kernel for eagle (#15909)
Ekagra Ranjan committed
Apr 1, 2025 at 21:15 UTC
24b7fb455a77518f50b21c35a693ae28a5b49b0b
1 file changed
+1
-2
vllm/v1/spec_decode/eagle.py
+1
-2
@@ -250,13 +250,12 @@ def prepare_input_kernel(
250
num_tokens = end_pos - start_pos
251
252
index_start = tl.load(cu_query_lens_ptr + pid)
253
- indices = index_start + tl.arange(0, BLOCK_SIZE)
253
254
num_blocks = tl.cdiv(num_tokens, BLOCK_SIZE)
255
for i in tl.range(num_blocks):
256
offset = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
257
tl.store(
258
out_ptr + start_pos + offset,
260
- indices,
259
+ index_start + offset,
260
mask=offset < num_tokens,
261
)