125 lines
3.9 KiB
Python
125 lines
3.9 KiB
Python
"""
|
|
Test D2H transfer performance with pinned vs non-contiguous memory.
|
|
"""
|
|
|
|
import torch
|
|
import time
|
|
|
|
print("=" * 60)
|
|
print("Test: D2H Transfer Performance (for nsys profiling)")
|
|
print("=" * 60)
|
|
|
|
# Setup
|
|
num_layers = 8
|
|
num_blocks = 16
|
|
block_size = 1024
|
|
num_kv_heads = 8
|
|
head_dim = 64
|
|
|
|
# Allocate CPU cache (pinned)
|
|
k_cache_cpu = torch.zeros(
|
|
num_layers, num_blocks, block_size, num_kv_heads, head_dim,
|
|
dtype=torch.float16, device="cpu", pin_memory=True
|
|
)
|
|
|
|
# Allocate GPU cache
|
|
k_cache_gpu = torch.randn(
|
|
num_layers, 4, block_size, num_kv_heads, head_dim,
|
|
dtype=torch.float16, device="cuda"
|
|
)
|
|
|
|
# Warmup
|
|
print("\nWarmup...")
|
|
for _ in range(10):
|
|
k_cache_cpu[:, 0].copy_(k_cache_gpu[:, 0], non_blocking=True)
|
|
torch.cuda.synchronize()
|
|
|
|
print(f"\nTensor info:")
|
|
print(f" k_cache_cpu.is_pinned(): {k_cache_cpu.is_pinned()}")
|
|
print(f" k_cache_cpu.is_contiguous(): {k_cache_cpu.is_contiguous()}")
|
|
print(f" k_cache_cpu[:, 0].is_pinned(): {k_cache_cpu[:, 0].is_pinned()}")
|
|
print(f" k_cache_cpu[:, 0].is_contiguous(): {k_cache_cpu[:, 0].is_contiguous()}")
|
|
|
|
# Test 1: Non-contiguous slice (current approach)
|
|
print(f"\n" + "=" * 60)
|
|
print("Test 1: Non-contiguous slice copy (current approach)")
|
|
print("=" * 60)
|
|
|
|
NUM_ITERS = 50 # Reduced for profiling
|
|
|
|
torch.cuda.nvtx.range_push("Test1_NonContiguous")
|
|
times = []
|
|
for i in range(NUM_ITERS):
|
|
torch.cuda.nvtx.range_push(f"D2H_NonContig_{i}")
|
|
start = time.perf_counter()
|
|
k_cache_cpu[:, i % num_blocks].copy_(k_cache_gpu[:, 0], non_blocking=True)
|
|
torch.cuda.synchronize()
|
|
times.append(time.perf_counter() - start)
|
|
torch.cuda.nvtx.range_pop()
|
|
torch.cuda.nvtx.range_pop()
|
|
|
|
avg_time = sum(times) / len(times)
|
|
print(f"Average time: {avg_time * 1000:.3f} ms")
|
|
print(f"Bandwidth: {k_cache_gpu[:, 0].numel() * 2 / avg_time / 1e9:.2f} GB/s")
|
|
|
|
# Test 2: Transpose to make dimension contiguous
|
|
print(f"\n" + "=" * 60)
|
|
print("Test 2: Transpose to contiguous dimension")
|
|
print("=" * 60)
|
|
|
|
# Reshape to [num_blocks, num_layers, block_size, num_kv_heads, head_dim]
|
|
k_cache_cpu_transposed = torch.zeros(
|
|
num_blocks, num_layers, block_size, num_kv_heads, head_dim,
|
|
dtype=torch.float16, device="cpu", pin_memory=True
|
|
)
|
|
|
|
print(f" k_cache_cpu_transposed[0].is_pinned(): {k_cache_cpu_transposed[0].is_pinned()}")
|
|
print(f" k_cache_cpu_transposed[0].is_contiguous(): {k_cache_cpu_transposed[0].is_contiguous()}")
|
|
|
|
torch.cuda.nvtx.range_push("Test2_Contiguous")
|
|
times = []
|
|
for i in range(NUM_ITERS):
|
|
torch.cuda.nvtx.range_push(f"D2H_Contig_{i}")
|
|
start = time.perf_counter()
|
|
k_cache_cpu_transposed[i % num_blocks].copy_(k_cache_gpu[:, 0], non_blocking=True)
|
|
torch.cuda.synchronize()
|
|
times.append(time.perf_counter() - start)
|
|
torch.cuda.nvtx.range_pop()
|
|
torch.cuda.nvtx.range_pop()
|
|
|
|
avg_time = sum(times) / len(times)
|
|
print(f"Average time: {avg_time * 1000:.3f} ms")
|
|
print(f"Bandwidth: {k_cache_gpu[:, 0].numel() * 2 / avg_time / 1e9:.2f} GB/s")
|
|
|
|
# Test 3: Fully contiguous buffer
|
|
print(f"\n" + "=" * 60)
|
|
print("Test 3: Fully contiguous buffer")
|
|
print("=" * 60)
|
|
|
|
k_cache_cpu_flat = torch.zeros(
|
|
num_layers * block_size * num_kv_heads * head_dim,
|
|
dtype=torch.float16, device="cpu", pin_memory=True
|
|
)
|
|
|
|
print(f" k_cache_cpu_flat.is_pinned(): {k_cache_cpu_flat.is_pinned()}")
|
|
print(f" k_cache_cpu_flat.is_contiguous(): {k_cache_cpu_flat.is_contiguous()}")
|
|
|
|
torch.cuda.nvtx.range_push("Test3_FlatContiguous")
|
|
times = []
|
|
for i in range(NUM_ITERS):
|
|
torch.cuda.nvtx.range_push(f"D2H_Flat_{i}")
|
|
start = time.perf_counter()
|
|
k_cache_cpu_flat.copy_(k_cache_gpu[:, 0].flatten(), non_blocking=True)
|
|
torch.cuda.synchronize()
|
|
times.append(time.perf_counter() - start)
|
|
torch.cuda.nvtx.range_pop()
|
|
torch.cuda.nvtx.range_pop()
|
|
|
|
avg_time = sum(times) / len(times)
|
|
print(f"Average time: {avg_time * 1000:.3f} ms")
|
|
print(f"Bandwidth: {k_cache_cpu_flat.numel() * 2 / avg_time / 1e9:.2f} GB/s")
|
|
|
|
print("\n" + "=" * 60)
|
|
print("test_pinned_transfer: PASSED")
|
|
print("=" * 60)
|