[WIP] fixing attention compute error.

nanovllm/engine/llm_engine.py

@@ -31,6 +31,8 @@ class LLMEngine:
         self.model_runner = ModelRunner(config, 0, self.events)
         self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True)
         config.eos = self.tokenizer.eos_token_id
+        # Set Sequence.block_size to match the KV cache block size
+        Sequence.block_size = config.kvcache_block_size
         self.scheduler = Scheduler(config, self.model_runner.kvcache_manager)
         atexit.register(self.exit)
 

nanovllm/engine/model_runner.py

@@ -521,6 +521,7 @@ class ModelRunner:
         print(f"[Ring Buffer Prefill] Complete: {chunk_idx} chunks", file=sys.stderr)
 
         # Sample from last logits
+        # For chunked prefill, ParallelLMHead automatically selects last position's logits
         temperatures = self.prepare_sample(seqs) if self.rank == 0 else None
         if logits is not None:
             token_ids = self.sampler(logits, temperatures).tolist() if self.rank == 0 else None