From a7465ba67d6ad0414737ff1db25ff7f59abd9dd8 Mon Sep 17 00:00:00 2001 From: OlivierDehaene <23298448+OlivierDehaene@users.noreply.github.com> Date: Fri, 25 Oct 2024 10:37:10 +0200 Subject: [PATCH] fix kernel --- server/text_generation_server/models/metadata_kernels.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/server/text_generation_server/models/metadata_kernels.py b/server/text_generation_server/models/metadata_kernels.py index 2238631b..2357564e 100644 --- a/server/text_generation_server/models/metadata_kernels.py +++ b/server/text_generation_server/models/metadata_kernels.py @@ -218,7 +218,7 @@ def triton_copy_next_input_ids_inplace( mask = (next_input_ids_start + block_arange) < next_input_ids_end # Mask values for request still prefilling - decode_mask = (cache_length + input_length + block_arange) > prompt_length + decode_mask = (cache_length + input_length + block_arange) >= prompt_length mask = mask & decode_mask @@ -229,7 +229,11 @@ def triton_copy_next_input_ids_inplace( # Store in all_input_ids, since it is a 2D tensor, apply stride * bid tl.store( - all_input_ids_ptr + stride_all_input_ids * bid + cache_length + block_arange, + all_input_ids_ptr + + stride_all_input_ids * bid + + cache_length + + input_length + + block_arange, next_input_ids, mask=mask, )