From fd92054e1d89c9d89fa5673def47b29b21055071 Mon Sep 17 00:00:00 2001 From: Curtis Ruck Date: Fri, 11 Apr 2025 10:09:24 -0400 Subject: [PATCH] Fix state.plan call to use positional arguments --- .../layers/attention/flashinfer.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/server/text_generation_server/layers/attention/flashinfer.py b/server/text_generation_server/layers/attention/flashinfer.py index 9479b606717..7095593b936 100644 --- a/server/text_generation_server/layers/attention/flashinfer.py +++ b/server/text_generation_server/layers/attention/flashinfer.py @@ -84,16 +84,16 @@ def use_prefill_with_paged_kv_state( token = prefill_with_paged_kv_state.set(state) try: state.plan( - qo_indptr=cu_seqlens, - paged_kv_indptr=indptr, - paged_kv_indices=block_tables, - paged_kv_last_page_len=last_page_len, - num_qo_heads=num_heads, - num_kv_heads=num_kv_heads, - head_dim=head_size, - kv_data_type=kv_dtype, + cu_seqlens, + indptr, + block_tables, + last_page_len, + num_heads, + num_kv_heads, + head_size, + page_size, q_data_type=q_dtype, - page_size=page_size, + kv_data_type=kv_dtype, ) yield finally: