[WIP] fixing attention compute error.

2025-12-30 00:31:48 +08:00
parent bf4c63c7ec
commit 89f8020d38
12 changed files with 2175 additions and 103 deletions
--- a/nanovllm/engine/model_runner.py
+++ b/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