[fix] Fixed needle test bug.

2026-01-05 18:34:09 +08:00
parent d623043a3c
commit 054aaff403
5 changed files with 8 additions and 4 deletions
--- a/nanovllm/layers/attention.py
+++ b/nanovllm/layers/attention.py
@@ -84,6 +84,10 @@ class Attention(nn.Module):
            context.kvcache_manager is not None and
            hasattr(context.kvcache_manager, 'offload_engine')
        )
+        
+        #! Ensure synchronization before accessing k_cache/v_cache
+        torch.cuda.synchronize()
+        #! =======================================================

        if is_chunked_offload:
            # Chunked offload mode: use compute_stream for store_kvcache