diff --git a/bench_offload.py b/bench_offload.py index 97224aa..33d46cf 100644 --- a/bench_offload.py +++ b/bench_offload.py @@ -5,7 +5,7 @@ from nanovllm import LLM, SamplingParams def bench_decode(llm, num_seqs, max_input_len, max_output_len): - """Benchmark decode performance""" + """Benchmark decode performance (original test)""" seed(0) prompt_token_ids = [[randint(0, 10000) for _ in range(randint(100, max_input_len))] for _ in range(num_seqs)] sampling_params = [SamplingParams(temperature=0.6, ignore_eos=True, max_tokens=randint(100, max_output_len)) for _ in range(num_seqs)] @@ -21,6 +21,7 @@ def bench_decode(llm, num_seqs, max_input_len, max_output_len): def bench_prefill(llm, num_seqs, input_len): """Benchmark prefill performance""" seed(0) + # Fixed length input, minimal output to focus on prefill prompt_token_ids = [[randint(0, 10000) for _ in range(input_len)] for _ in range(num_seqs)] sampling_params = SamplingParams(temperature=0.6, ignore_eos=True, max_tokens=1) @@ -40,6 +41,8 @@ def main(): max_model_len=128 * 1024, max_num_batched_tokens=128 * 1024, enable_cpu_offload=True, + num_gpu_blocks=6, + num_prefetch_blocks=2, ) # Warmup @@ -48,15 +51,16 @@ def main(): print("=" * 60) print("Prefill Benchmark (CPU Offload)") print("=" * 60) - bench_prefill(llm, num_seqs=1, input_len=64*1024) - # bench_prefill(llm, num_seqs=1, input_len=16384) - # bench_prefill(llm, num_seqs=1, input_len=32000) + # bench_prefill(llm, num_seqs=1, input_len=1024) + # bench_prefill(llm, num_seqs=1, input_len=2048) + # bench_prefill(llm, num_seqs=1, input_len=4096) + bench_prefill(llm, num_seqs=1, input_len=8192) print("=" * 60) print("Decode Benchmark (CPU Offload)") print("=" * 60) - bench_decode(llm, num_seqs=1, max_input_len=64*1024, max_output_len=256) - # bench_decode(llm, num_seqs=1, max_input_len=16384, max_output_len=256) + bench_decode(llm, num_seqs=1, max_input_len=1024, max_output_len=128) + # bench_decode(llm, num_seqs=1, max_input_len=2048, max_output_len=128) if __name__ == "__main__": diff --git a/nanovllm/engine/model_runner.py b/nanovllm/engine/model_runner.py index e7f3c8a..caab2ed 100644 --- a/nanovllm/engine/model_runner.py +++ b/nanovllm/engine/model_runner.py @@ -7,7 +7,7 @@ from multiprocessing.shared_memory import SharedMemory from nanovllm.config import Config from nanovllm.engine.sequence import Sequence from nanovllm.models.qwen3 import Qwen3ForCausalLM -from nanovllm.layers.sampler import Sampler +from nanovllm.layers.sampler import GreedySampler from nanovllm.utils.context import set_context, get_context, reset_context from nanovllm.utils.loader import load_model from nanovllm.utils.logger import get_logger @@ -34,7 +34,7 @@ class ModelRunner: torch.set_default_device("cuda") self.model = Qwen3ForCausalLM(hf_config) load_model(self.model, config.model) - self.sampler = Sampler() + self.sampler = GreedySampler() self.warmup_model() self.allocate_kv_cache() if not self.enforce_eager: diff --git a/nanovllm/kvcache/hybrid_manager.py b/nanovllm/kvcache/hybrid_manager.py index 019c27b..cc02f8b 100644 --- a/nanovllm/kvcache/hybrid_manager.py +++ b/nanovllm/kvcache/hybrid_manager.py @@ -1039,6 +1039,8 @@ class HybridKVCacheManager(KVCacheManager): """ assert not seq.block_table, "Sequence already has blocks" + h = -1 # Running hash for prefix cache + for i in range(seq.num_blocks): # Allocate CPU block if not self.free_cpu_blocks: @@ -1049,10 +1051,19 @@ class HybridKVCacheManager(KVCacheManager): cpu_block_id = self.free_cpu_blocks.popleft() + # Get token IDs for this block and compute hash + token_ids = seq.block(i) + if len(token_ids) == self._block_size: + h = self.compute_hash(token_ids, h) + else: + h = -1 # Incomplete block + # Allocate logical block logical_id = self.free_logical_ids.popleft() block = self.logical_blocks[logical_id] block.ref_count = 1 + block.hash = h + block.token_ids = token_ids.copy() if len(token_ids) == self._block_size else [] block.location = BlockLocation.CPU block.cpu_block_id = cpu_block_id block.gpu_slot = -1 @@ -1060,6 +1071,10 @@ class HybridKVCacheManager(KVCacheManager): self.cpu_block_to_logical[cpu_block_id] = logical_id seq.block_table.append(logical_id) + # Update prefix cache + if h != -1: + self.hash_to_logical_id[h] = logical_id + def get_cpu_block_table(self, seq: Sequence) -> List[int]: """ Get CPU block ID list for sequence. diff --git a/nanovllm/layers/sampler.py b/nanovllm/layers/sampler.py index b101018..dacb978 100644 --- a/nanovllm/layers/sampler.py +++ b/nanovllm/layers/sampler.py @@ -13,3 +13,13 @@ class Sampler(nn.Module): probs = torch.softmax(logits, dim=-1) sample_tokens = probs.div_(torch.empty_like(probs).exponential_(1).clamp_min_(1e-10)).argmax(dim=-1) return sample_tokens + + +class GreedySampler(nn.Module): + + def __init__(self): + super().__init__() + + @torch.compile + def forward(self, logits: torch.Tensor, temperatures: torch.Tensor = None): + return logits.argmax(dim=-1)