Commit81e9d5c

committed

rename chunked_mode to chunked_logits

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

1 parentb20f62d commit81e9d5cCopy full SHA for 81e9d5c

File tree

2 files changed

+15

-14

lines changed

tensorrt_llm/_torch/pyexecutor
- handle_logits.py
- llm_request.py

2 files changed

+15

-14

lines changed

`‎tensorrt_llm/_torch/pyexecutor/handle_logits.py‎`

Lines changed: 5 additions & 4 deletions

Original file line number	Diff line number	Diff line change
`@@ -79,8 +79,9 @@ def __call__(`
`79`	`79`	`logits_view=logits[logits_begin:logits_end].reshape(`
`80`	`80`	`1,beam_width,-1)`
`81`	`81`	`llm_req.py_result.append_generation_logits(logits_view)`
`82`		`-`
`83`		`-# Finalize any remaining logits transfers for all requests(chunked mode)`
	`82`	`+`
	`83`	`+# Finalize any remaining logits transfers for all requestsinchunked mode`
`84`	`84`	`forllm_reqinchain(context_requests,generation_requests):`
`85`		`-ifllm_req.py_return_generation_logitsorllm_req.py_return_context_logits:`
`86`		`-llm_req.py_result.post_processing_transfer()`
	`85`	`+ifllm_req.py_use_chunked_logits:`
	`86`	`+ifllm_req.py_return_generation_logitsorllm_req.py_return_context_logits:`
	`87`	`+llm_req.py_result.post_processing_transfer()`

`‎tensorrt_llm/_torch/pyexecutor/llm_request.py‎`

Lines changed: 10 additions & 10 deletions

Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@ def __init__(self,`
`47`	`47`	`seq_length:int,`
`48`	`48`	`use_device_memory=True,`
`49`	`49`	`should_exclude_last=False,`
`50`		`-chunked_mode=False,`
	`50`	`+use_chunked_logits=False,`
`51`	`51`	`streaming=False,`
`52`	`52`	`chunk_size=8):`
`53`	`53`	`ifshould_exclude_last:`
`@@ -57,7 +57,7 @@ def __init__(self,`
`57`	`57`	`self.seq_length=seq_length`
`58`	`58`	`self.use_device_memory=use_device_memory`
`59`	`59`	`self._should_exclude_last=should_exclude_last`
`60`		`-self.chunked_mode=chunked_mode`
	`60`	`+self.use_chunked_logits=use_chunked_logits`
`61`	`61`	`self.chunk_size=chunk_size`
`62`	`62`	`self.streaming=streaming`
`63`	`63`	`self._logits_indices= []`
`@@ -68,7 +68,7 @@ def __init__(self,`
`68`	`68`	`self.vocab_size=-1`
`69`	`69`
`70`	`70`	`# Chunked mode: device-side fragments`
`71`		`-ifchunked_mode:`
	`71`	`+ifuse_chunked_logits:`
`72`	`72`	`self._device_fragments:List[torch.Tensor]= []`
`73`	`73`	`self._current_position=0`
`74`	`74`
`@@ -105,7 +105,7 @@ def append(self, logits: torch.Tensor):`
`105`	`105`	`logits=logits.unsqueeze(1)`
`106`	`106`	`assertlogits.ndim==3,f"Bad logits shape, expect [num_tokens, beam_width, vocab_size], got{logits.shape}"`
`107`	`107`
`108`		`-ifself.chunked_mode:`
	`108`	`+ifself.use_chunked_logits:`
`109`	`109`	`ifself.beam_width==-1:`
`110`	`110`	`self._init_chunked_storage(logits)`
`111`	`111`	`self._add_fragment(logits)`
`@@ -181,7 +181,7 @@ def _transfer_chunk_to_host(self):`
`181`	`181`
`182`	`182`	`deffinalize_transfer(self):`
`183`	`183`	`"""Force transfer of any remaining fragments to host (for chunked mode)"""`
`184`		`-ifself.chunked_modeandhasattr(`
	`184`	`+ifself.use_chunked_logitsandhasattr(`
`185`	`185`	`self,'_device_fragments')andself._device_fragments:`
`186`	`186`	`self._transfer_chunk_to_host()`
`187`	`187`
`@@ -243,20 +243,20 @@ def __init__(self,`
`243`	`243`	`return_context_logits:bool=False,`
`244`	`244`	`return_generation_logits:bool=False,`
`245`	`245`	`exclude_last_generation_logits:bool=False,`
`246`		`-chunked_mode:bool=False,`
	`246`	`+use_chunked_logits:bool=True,`
`247`	`247`	`chunk_size:int=8):`
`248`	`248`	`self._streaming=streaming`
`249`	`249`	`self._context_logits=LogitsStorage(`
`250`	`250`	`prompt_len,`
`251`	`251`	`use_device_memory,`
`252`		`-chunked_mode=chunked_mode,`
	`252`	`+use_chunked_logits=use_chunked_logits,`
`253`	`253`	`streaming=streaming,`
`254`	`254`	`chunk_size=chunk_size)ifreturn_context_logitselseNone`
`255`	`255`	`self._generation_logits=LogitsStorage(`
`256`	`256`	`max_new_tokens,`
`257`	`257`	`use_device_memory,`
`258`	`258`	`exclude_last_generation_logits,`
`259`		`-chunked_mode=chunked_mode,`
	`259`	`+use_chunked_logits=use_chunked_logits,`
`260`	`260`	`streaming=streaming,`
`261`	`261`	`chunk_size=chunk_size)ifreturn_generation_logitselseNone`
`262`	`262`	`self._log_probs=LogProbStorage()ifreturn_log_probselseNone`
`@@ -394,7 +394,7 @@ def __init__(`
`394`	`394`	`is_draft:bool=False,`
`395`	`395`	`seq_slot:Optional[int]=None,`
`396`	`396`	`target_seq_slot:Optional[int]=None,`
`397`		`-use_chunked_logits:bool=False,`
	`397`	`+use_chunked_logits:bool=True,`
`398`	`398`	`logits_chunk_size:int=8,`
`399`	`399`	`**kwargs):`
`400`	`400`
`@@ -466,7 +466,7 @@ def __init__(`
`466`	`466`	`return_context_logits,`
`467`	`467`	`return_generation_logits,`
`468`	`468`	`exclude_last_generation_logits,`
`469`		`-chunked_mode=use_chunked_logits,`
	`469`	`+use_chunked_logits=use_chunked_logits,`
`470`	`470`	`chunk_size=logits_chunk_size)`
`471`	`471`	`self.child_requests= []`
`472`	`472`

0 commit comments

Comments

(0)

Movatterモバイル変換

Navigation Menu

Search code, repositories, users, issues, pull requests...

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Commit81e9d5c

File tree

2 files changed

2 files changed

`‎tensorrt_llm/_torch/pyexecutor/handle_logits.py‎`

`‎tensorrt_llm/_torch/pyexecutor/llm_request.py‎`

0 commit comments