fix kernel

2025-09-11 20:34:54 +00:00 · 2024-10-24 19:17:31 +02:00 · 2024-10-24 19:17:31 +02:00 · 347f3f51da
commit 347f3f51da
parent d1e95ceaff
1 changed files with 1 additions and 5 deletions
--- a/server/text_generation_server/models/metadata_kernels.py
+++ b/server/text_generation_server/models/metadata_kernels.py
@ -229,11 +229,7 @@ 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
-        + input_length
-        + block_arange,
+        all_input_ids_ptr + stride_all_input_ids * bid + cache_length + block_arange,
        next_input_ids,
        mask=mask,
    )