zijie-tian/nano-vllm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[refactor] Before add sprae policy.

Browse Source
This commit is contained in:
Zijie Tian
2026-01-05 21:19:24 +08:00
parent e554d5482b
commit c7ac39dfbd
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
nanovllm/engine/model_runner.py
View File

@@ -382,7 +382,7 @@ class ModelRunner:
return self.model.compute_logits(graph_vars["outputs"][:bs]) return self.model.compute_logits(graph_vars["outputs"][:bs])
def run(self, seqs: list[Sequence], is_prefill: bool) -> list[int]: def run(self, seqs: list[Sequence], is_prefill: bool) -> list[int]:
# Check if Chunked Offload mode should be used (all blocks on CPU) #> Check if Chunked Offload mode should be used (all blocks on CPU)
if hasattr(self, 'kvcache_manager') and hasattr(self.kvcache_manager, 'get_all_cpu_blocks'): if hasattr(self, 'kvcache_manager') and hasattr(self.kvcache_manager, 'get_all_cpu_blocks'):
use_chunked_offload = self._should_use_chunked_offload(seqs, is_prefill) use_chunked_offload = self._should_use_chunked_offload(seqs, is_prefill)
if use_chunked_offload: if use_chunked_offload:
@@ -391,6 +391,7 @@ class ModelRunner:
else: else:
return self.run_chunked_offload_decode(seqs) return self.run_chunked_offload_decode(seqs)
#> Following Code will not use Chunked Offload mode
input_ids, positions = self.prepare_prefill(seqs) if is_prefill else self.prepare_decode(seqs) input_ids, positions = self.prepare_prefill(seqs) if is_prefill else self.prepare_decode(seqs)
temperatures = self.prepare_sample(seqs) if self.rank == 0 else None temperatures = self.prepare_sample(seqs) if self.rank == 0 else None
logits = self.run_model(input_ids, positions, is_prefill) logits = self.run_model(input_ids, positions, is_prefill)