[WIP] NEED to modify communication.

tests/test_pinned_memory_slice.py

@@ -0,0 +1,70 @@
+"""
+Test if slicing maintains pinned memory property.
+"""
+
+import torch
+
+print("=" * 60)
+print("Test: Pinned Memory Property with Slicing")
+print("=" * 60)
+
+# Create a pinned tensor with shape similar to k_cache_cpu
+# [num_layers, num_cpu_blocks, block_size, num_kv_heads, head_dim]
+tensor = torch.zeros(8, 16, 1024, 8, 64, dtype=torch.float16, device="cpu", pin_memory=True)
+
+print(f"\n1. Original tensor:")
+print(f"   - Shape: {tensor.shape}")
+print(f"   - is_pinned(): {tensor.is_pinned()}")
+print(f"   - is_contiguous(): {tensor.is_contiguous()}")
+
+# Test slicing operation (what we do in offload_slot_to_cpu)
+slice_view = tensor[:, 0]  # Same as k_cache_cpu[:, cpu_block_id]
+
+print(f"\n2. Sliced tensor [:, 0]:")
+print(f"   - Shape: {slice_view.shape}")
+print(f"   - is_pinned(): {slice_view.is_pinned()}")
+print(f"   - is_contiguous(): {slice_view.is_contiguous()}")
+
+# Test if contiguous() helps
+contiguous_slice = tensor[:, 0].contiguous()
+
+print(f"\n3. Contiguous slice [:, 0].contiguous():")
+print(f"   - Shape: {contiguous_slice.shape}")
+print(f"   - is_pinned(): {contiguous_slice.is_pinned()}")
+print(f"   - is_contiguous(): {contiguous_slice.is_contiguous()}")
+
+# Test copy behavior
+gpu_tensor = torch.zeros(8, 4, 1024, 8, 64, dtype=torch.float16, device="cuda")
+gpu_slice = gpu_tensor[:, 0]
+
+print(f"\n4. GPU tensor slice:")
+print(f"   - Shape: {gpu_slice.shape}")
+print(f"   - is_contiguous(): {gpu_slice.is_contiguous()}")
+
+# Simulate the problematic copy operation
+print(f"\n5. Testing copy operations:")
+
+# Method 1: Direct slice copy (current approach - SLOW)
+slice_dst = tensor[:, 1]
+print(f"   Method 1 (slice view): dst.is_pinned()={slice_dst.is_pinned()}")
+
+# Method 2: Use contiguous destination
+contiguous_dst = tensor[:, 2].contiguous()
+print(f"   Method 2 (contiguous): dst.is_pinned()={contiguous_dst.is_pinned()}")
+
+print("\n" + "=" * 60)
+print("Conclusion:")
+print("=" * 60)
+
+if not slice_view.is_pinned():
+    print("❌ Slicing LOSES pinned memory property!")
+    print("   This causes Device-to-Pageable transfers (SLOW)")
+else:
+    print("✓ Slicing maintains pinned memory property")
+
+if contiguous_slice.is_pinned():
+    print("✓ .contiguous() maintains pinned memory property")
+else:
+    print("❌ .contiguous() also loses pinned memory property")
+
+print("\n" + "=" * 60)