From 4bf0da3e5c8bb4a920efcb9de014443a7cd6327a Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Tue, 25 Jun 2024 00:57:56 +0000 Subject: [PATCH] format fix --- examples/offline_profile.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/examples/offline_profile.py b/examples/offline_profile.py index 631248246f1f7..959431c9d39cb 100644 --- a/examples/offline_profile.py +++ b/examples/offline_profile.py @@ -75,16 +75,15 @@ def run_profile(context: ProfileContext, csv_output: Optional[str], sys.exit(-1) for i in range(batch_size): - prompt_token_ids=torch.randint( - 128, # 128 to skip over special tokens - llm.llm_engine.model_config.get_vocab_size() // 2, - size=(prompt_len, )).tolist() - prompt_token_ids = {'prompt_token_ids' : prompt_token_ids} - - llm.llm_engine.add_request( - request_id=f"seq{i}", - inputs=prompt_token_ids, - params=sampling_params) + prompt_token_ids = torch.randint( + 128, # 128 to skip over special tokens + llm.llm_engine.model_config.get_vocab_size() // 2, + size=(prompt_len, )).tolist() + prompt_token_ids = {'prompt_token_ids': prompt_token_ids} + + llm.llm_engine.add_request(request_id=f"seq{i}", + inputs=prompt_token_ids, + params=sampling_params) with nm_profile() as prefill_prof: llm.llm_engine.step() # First step is prefill