zijie-tian/nano-vllm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[WIP] Before modify to FlashInfer.

Browse Source
This commit is contained in:
Zijie Tian
2025-12-30 01:11:13 +08:00
parent 89f8020d38
commit 7af721c12c
4 changed files with 104 additions and 480 deletions
Download Patch File Download Diff File Expand all files Collapse all files

104
tests/test_flashinfer_merge.py Normal file
View File

@@ -0,0 +1,104 @@
"""
Test FlashInfer chunked attention with CPU offload.
Uses single_prefill_with_kv_cache + merge_state for chunked KV processing.
"""
import torch
import flashinfer
# ============================================================
# Core Functions
# ============================================================
def chunked_prefill_causal(q, k_cpu, v_cpu, q_chunk_size, kv_chunk_size):
"""
Chunked causal attention with KV on CPU.
q: [seq_q, num_heads, head_dim] on GPU
k_cpu, v_cpu: [seq_kv, num_kv_heads, head_dim] on CPU
"""
seq_q = q.shape[0]
seq_kv = k_cpu.shape[0]
final_outputs = []
for q_start in range(0, seq_q, q_chunk_size):
q_end = min(q_start + q_chunk_size, seq_q)
q_chunk = q[q_start:q_end]
merged_output = None
merged_lse = None
for kv_start in range(0, seq_kv, kv_chunk_size):
kv_end = min(kv_start + kv_chunk_size, seq_kv)
if kv_start >= q_end:
continue
k_chunk = k_cpu[kv_start:kv_end].to(q.device, non_blocking=True)
v_chunk = v_cpu[kv_start:kv_end].to(q.device, non_blocking=True)
causal = not (kv_end <= q_start)
partial_out, partial_lse = flashinfer.single_prefill_with_kv_cache(
q_chunk, k_chunk, v_chunk,
causal=causal,
return_lse=True,
)
if merged_output is None:
merged_output, merged_lse = partial_out, partial_lse
else:
merged_output, merged_lse = flashinfer.merge_state(
merged_output, merged_lse,
partial_out, partial_lse,
)
final_outputs.append(merged_output)
return torch.cat(final_outputs, dim=0)
# ============================================================
# Main Test Script
# ============================================================
print("=" * 60)
print("Testing FlashInfer chunked attention with CPU offload")
print("=" * 60)
num_heads = 32
num_kv_heads = 8
head_dim = 128
test_configs = [
(32768, 8192, 8192), # 32K tokens
(65536, 8192, 8192), # 64K tokens
(131072, 16384, 16384), # 128K tokens
# (262144, 16384, 16384), # 256K tokens (slow)
# (524288, 16384, 16384), # 512K tokens (slow)
]
for seq_len, q_chunk, kv_chunk in test_configs:
torch.manual_seed(42)
q = torch.randn(seq_len, num_heads, head_dim, dtype=torch.float16, device='cuda')
k_cpu = torch.randn(seq_len, num_kv_heads, head_dim, dtype=torch.float16, device='cpu')
v_cpu = torch.randn(seq_len, num_kv_heads, head_dim, dtype=torch.float16, device='cpu')
# Chunked result
chunked_out = chunked_prefill_causal(q, k_cpu, v_cpu, q_chunk, kv_chunk)
# Reference
k_gpu = k_cpu.to('cuda')
v_gpu = v_cpu.to('cuda')
ref_out = flashinfer.single_prefill_with_kv_cache(q, k_gpu, v_gpu, causal=True)
max_diff = (ref_out - chunked_out).abs().max().item()
mean_diff = (ref_out - chunked_out).abs().mean().item()
num_chunks = (seq_len + q_chunk - 1) // q_chunk
assert max_diff < 1e-2, f"FAILED: max_diff={max_diff:.6f}"
print(f"seq={seq_len//1024}K, chunks={num_chunks}: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}")
print("\ntest_flashinfer_merge: PASSED")

70
tests/test_pinned_memory_slice.py
View File

@@ -1,70 +0,0 @@
"""
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)

124
tests/test_pinned_transfer.py
View File

@@ -1,124 +0,0 @@
"""
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)

286
tests/test_sim.py
View File

@@ -1,286 +0,0 @@
"""
Chunked Prefill + KV Cache Offload Simulation v2
改进：
1. 简化日志输出
2. 添加reduce时间
3. 计算必须等待KV load完成
"""
import threading
import time
from dataclasses import dataclass
from typing import Optional
from concurrent.futures import ThreadPoolExecutor, Future
# ============== 配置参数 ==============
NUM_CHUNKS = 8
GPU_SLOTS = 4
# 模拟时间 (秒)
TIME_COMPUTE_BLOCK = 0.10 # 计算一个attention block
TIME_REDUCE = 0.03 # 两个partial result做一次reduce
TIME_TRANSFER = 0.08 # 传输一个KV cache
TIME_PROJ = 0.02 # projection生成KV
# ============== 全局时间基准 ==============
START_TIME = None
def now() -> float:
"""返回相对开始的时间(ms)"""
return (time.time() - START_TIME) * 1000
def log_compute(msg: str):
"""计算队列日志（无缩进）"""
print(f"[{now():7.1f}ms] [COMPUTE] {msg}")
def log_transfer(msg: str):
"""传输队列日志（缩进）"""
print(f"[{now():7.1f}ms] [TRANSFER] {msg}")
def log_info(msg: str):
"""一般信息"""
print(f"[{now():7.1f}ms] {msg}")
# ============== GPU Slot管理 ==============
class GPUSlots:
def __init__(self, num_slots: int):
self.slots = [None] * num_slots # slot_id -> kv_idx
self.kv_to_slot = {} # kv_idx -> slot_id
self.lock = threading.Lock()
# KV ready events: kv_idx -> Event
self.kv_ready = {}
def alloc(self, kv_idx: int) -> int:
with self.lock:
for sid, val in enumerate(self.slots):
if val is None:
self.slots[sid] = kv_idx
self.kv_to_slot[kv_idx] = sid
# 创建ready event
if kv_idx not in self.kv_ready:
self.kv_ready[kv_idx] = threading.Event()
return sid
raise RuntimeError(f"No free slot for KV{kv_idx}")
def free(self, slot_id: int):
with self.lock:
kv_idx = self.slots[slot_id]
if kv_idx is not None:
del self.kv_to_slot[kv_idx]
# 清除event
if kv_idx in self.kv_ready:
del self.kv_ready[kv_idx]
self.slots[slot_id] = None
def free_kv(self, kv_idx: int):
with self.lock:
if kv_idx in self.kv_to_slot:
sid = self.kv_to_slot[kv_idx]
self.slots[sid] = None
del self.kv_to_slot[kv_idx]
if kv_idx in self.kv_ready:
del self.kv_ready[kv_idx]
def mark_ready(self, kv_idx: int):
"""标记KV已就绪（load完成或proj完成）"""
with self.lock:
if kv_idx in self.kv_ready:
self.kv_ready[kv_idx].set()
def wait_ready(self, kv_idx: int):
"""等待KV就绪"""
with self.lock:
event = self.kv_ready.get(kv_idx)
if event:
event.wait()
def has_kv(self, kv_idx: int) -> bool:
with self.lock:
return kv_idx in self.kv_to_slot
def state(self) -> str:
with self.lock:
return "[" + "][".join(
f"KV{v}" if v is not None else "----"
for v in self.slots
) + "]"
# ============== 操作执行 ==============
class Executor:
def __init__(self, gpu: GPUSlots):
self.gpu = gpu
self.compute_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="Compute")
self.transfer_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="Transfer")
def proj_kv(self, q_idx: int) -> Future:
"""Projection生成KV，返回Future"""
def task():
log_compute(f"PROJ Q{q_idx}->KV{q_idx} START")
time.sleep(TIME_PROJ)
slot_id = self.gpu.alloc(q_idx)
self.gpu.mark_ready(q_idx) # proj完成，KV立即可用
log_compute(f"PROJ Q{q_idx}->KV{q_idx} END slot={slot_id} | {self.gpu.state()}")
return slot_id
return self.compute_pool.submit(task)
def compute_attn(self, q_idx: int, kv_indices: list) -> Future:
"""计算attention block，会等待所有KV就绪"""
def task():
# 等待所有需要的KV就绪
for kv_idx in kv_indices:
self.gpu.wait_ready(kv_idx)
kv_str = ",".join(map(str, kv_indices))
log_compute(f"ATTN Q{q_idx}*KV[{kv_str}] START")
time.sleep(TIME_COMPUTE_BLOCK * len(kv_indices))
log_compute(f"ATTN Q{q_idx}*KV[{kv_str}] END")
return (q_idx, kv_indices)
return self.compute_pool.submit(task)
def reduce(self, q_idx: int, num_partials: int) -> Future:
"""Online softmax reduce多个partial结果"""
def task():
if num_partials <= 1:
return
# n个partial需要n-1次两两reduce
num_reduces = num_partials - 1
log_compute(f"REDUCE Q{q_idx} ({num_partials} partials) START")
time.sleep(TIME_REDUCE * num_reduces)
log_compute(f"REDUCE Q{q_idx} END")
return self.compute_pool.submit(task)
def load_kv(self, kv_idx: int) -> Future:
"""从CPU load KV到GPU"""
def task():
if self.gpu.has_kv(kv_idx):
log_transfer(f"LOAD KV{kv_idx} SKIP (already on GPU)")
return kv_idx
slot_id = self.gpu.alloc(kv_idx)
log_transfer(f"LOAD KV{kv_idx} START -> slot{slot_id}")
time.sleep(TIME_TRANSFER)
self.gpu.mark_ready(kv_idx) # load完成，标记就绪
log_transfer(f"LOAD KV{kv_idx} END | {self.gpu.state()}")
return kv_idx
return self.transfer_pool.submit(task)
def offload_kv(self, kv_idx: int) -> Future:
"""从GPU offload KV到CPU"""
def task():
log_transfer(f"OFFLOAD KV{kv_idx} START")
time.sleep(TIME_TRANSFER)
self.gpu.free_kv(kv_idx)
log_transfer(f"OFFLOAD KV{kv_idx} END | {self.gpu.state()}")
return kv_idx
return self.transfer_pool.submit(task)
def shutdown(self):
self.compute_pool.shutdown(wait=True)
self.transfer_pool.shutdown(wait=True)
# ============== 调度器 ==============
def schedule_query(exe: Executor, q_idx: int):
"""调度单个Query的处理"""
print(f"\n{'='*50}")
log_info(f"===== Query {q_idx} START =====")
hist_kv = list(range(q_idx)) # 历史KV: 0 ~ q_idx-1
num_partials = 0
# Phase 1: Projection生成当前KV
proj_fut = exe.proj_kv(q_idx)
proj_fut.result() # 等待完成
# Phase 2: 对角块计算 + 同时prefetch历史KV
# 启动对角块计算
diag_fut = exe.compute_attn(q_idx, [q_idx])
num_partials += 1
# 同时prefetch历史KV (最多3个slot可用)
prefetch_slots = min(len(hist_kv), GPU_SLOTS - 1)
prefetch_kv = hist_kv[:prefetch_slots]
prefetch_futs = [exe.load_kv(kv) for kv in prefetch_kv]
# 等待对角块完成
diag_fut.result()
# Phase 3: Offload当前KV释放slot
offload_fut = exe.offload_kv(q_idx)
# 等待prefetch完成，然后计算这批历史KV
for f in prefetch_futs:
f.result()
if prefetch_kv:
hist_fut = exe.compute_attn(q_idx, prefetch_kv)
num_partials += 1
else:
hist_fut = None
# 等待offload完成
offload_fut.result()
# Phase 4: 处理剩余历史KV
remaining_kv = hist_kv[prefetch_slots:]
computed_kv = prefetch_kv.copy()
while remaining_kv:
# 等待上一批计算完成
if hist_fut:
hist_fut.result()
# 释放已计算的KV
for kv in computed_kv:
exe.gpu.free_kv(kv)
# Load下一批
batch_size = min(len(remaining_kv), GPU_SLOTS)
batch_kv = remaining_kv[:batch_size]
remaining_kv = remaining_kv[batch_size:]
load_futs = [exe.load_kv(kv) for kv in batch_kv]
for f in load_futs:
f.result()
# 计算这批
hist_fut = exe.compute_attn(q_idx, batch_kv)
num_partials += 1
computed_kv = batch_kv
# 等待最后一批计算完成
if hist_fut:
hist_fut.result()
# 清理GPU
for kv in computed_kv:
exe.gpu.free_kv(kv)
# Phase 5: Reduce所有partial results
reduce_fut = exe.reduce(q_idx, num_partials)
reduce_fut.result()
log_info(f"===== Query {q_idx} END =====")
def main():
global START_TIME
START_TIME = time.time()
print("Chunked Prefill + KV Cache Offload Simulation v2")
print(f"Config: {NUM_CHUNKS} chunks, {GPU_SLOTS} GPU slots")
print(f"Time: compute={TIME_COMPUTE_BLOCK}s, transfer={TIME_TRANSFER}s, reduce={TIME_REDUCE}s")
gpu = GPUSlots(GPU_SLOTS)
exe = Executor(gpu)
try:
for q_idx in range(NUM_CHUNKS):
schedule_query(exe, q_idx)
print(f"\n{'='*50}")
log_info(f"ALL DONE! Total: {now():.1f}ms")
finally:
exe.shutdown()
if __name__ == "__main__":
main()