From 58488dd3252d21a6872b5fde976944d483b307dc Mon Sep 17 00:00:00 2001 From: Mike Iovine <6158008+mikeiovine@users.noreply.github.com> Date: Tue, 20 May 2025 14:21:53 -0700 Subject: [PATCH] [Do not land] Temp changes for testing chunked attention --- .../lm-eval-harness/lm_eval_tensorrt_llm.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py b/examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py index 89e4a293db81..8f5a52411345 100644 --- a/examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py +++ b/examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py @@ -74,7 +74,7 @@ def __init__( self.moe_expert_parallel_size = moe_expert_parallel_size self.moe_backend = moe_backend trt_kv_cache_config = TRT_KvCacheConfig(enable_block_reuse=False) - trt_kv_cache_config.free_gpu_memory_fraction = free_gpu_memory_fraction + trt_kv_cache_config.free_gpu_memory_fraction = 0.5 if max_tokens_kv_cache is not None: trt_kv_cache_config.max_tokens = max_tokens_kv_cache @@ -109,7 +109,8 @@ def __init__( model=model, tensor_parallel_size=tp, trust_remote_code=trust_remote_code, - enable_chunked_prefill=False, + enable_chunked_prefill=True, + max_num_tokens=2048, pytorch_backend_config=pytorch_config, tokenizer=self.tokenizer, kv_cache_config=trt_kv_cache_config, @@ -292,11 +293,11 @@ def _get_sp(gen_kwargs): # process the output of the request i r_out: RequestOutput = futures.pop(i).result() stop_words = future_stop_words.pop(i) - for word in stop_words: - txt = r_out.outputs[0].text - word_index = txt.find(word) - if word_index >= 0: - txt = txt[:word_index] + txt = r_out.outputs[0].text + # for word in stop_words: + # word_index = txt.find(word) + # if word_index >= 0: + # txt = txt[:word_index] results.append(txt) return results