""" Test: 验证 xattn_estimate 与 KV chunking kernels 的一致性 使用真实 KV cache 数据，对比: 1. xattn_estimate (高层 API) 2. 三阶段 KV chunking (softmax_compute_partial_stats + merge + normalize) 三阶段 KV chunking 流程: 1. softmax_compute_partial_stats: 计算每个 KV chunk 的 (m, l) 2. merge_softmax_stats: Host 端合并所有 chunks 的 stats 3. softmax_normalize_and_block_sum: 使用全局 stats 归一化 支持两种数据格式: 1. offload 模式保存: {"query", "key", "stride", "threshold", "density", "layer_id"} 2. GPU-only 模式保存: {"Q", "K", "chunk_size", "block_size", "stride", "threshold", "mask", "attn_sums", ...} Usage: # 使用 offload 模式数据 CUDA_VISIBLE_DEVICES=0 PYTHONPATH=/home/zijie/Code/nano-vllm:$PYTHONPATH \ python tests/test_xattn_estimate_alignment.py # 使用 GPU-only 模式数据 CUDA_VISIBLE_DEVICES=0 PYTHONPATH=/home/zijie/Code/nano-vllm:$PYTHONPATH \ python tests/test_xattn_estimate_alignment.py --gpuonly """ import sys sys.path.insert(0, "/home/zijie/Code/nano-vllm") import argparse import torch import math from nanovllm.ops.xattn import ( xattn_estimate, flat_group_gemm_fuse_reshape, softmax_compute_partial_stats, softmax_normalize_and_block_sum, merge_softmax_stats, find_blocks_chunked, ) # ============================================================ # 命令行参数 # ============================================================ parser = argparse.ArgumentParser() parser.add_argument("--gpuonly", action="store_true", help="使用 GPU-only 模式保存的数据") parser.add_argument("--data-file", type=str, default=None, help="数据文件路径") parser.add_argument("--chunk-size", type=int, default=None, help="覆盖 CHUNK_SIZE (用于测试不同分块大小)") args = parser.parse_args() # ============================================================ # 参数配置 # ============================================================ if args.gpuonly: DATA_FILE = args.data_file or "/home/zijie/Code/nano-vllm/results/mask_alignment/gpuonly_layer0.pt" else: DATA_FILE = args.data_file or "/home/zijie/Code/nano-vllm/results/kvcache/qkv_32485.pt" device = "cuda" # ============================================================ # Step 1: 加载真实数据 # ============================================================ print("=" * 60) print("Step 1: 加载真实 KV cache 数据") print("=" * 60) data = torch.load(DATA_FILE, map_location="cpu") # 检测数据格式并加载 if "Q" in data: # GPU-only 模式保存的格式 print(f"[INFO] 检测到 GPU-only 模式数据格式") Q = data["Q"].to(device) K = data["K"].to(device) BSA_BLOCK_SIZE = data.get("block_size", 128) CHUNK_SIZE = data.get("chunk_size", 4096) STRIDE = data.get("stride", 8) THRESHOLD = data.get("threshold", 0.9) if isinstance(THRESHOLD, torch.Tensor): THRESHOLD = THRESHOLD.item() # GPU-only 模式保存了 mask 和 attn_sums，可以用于验证 saved_mask = data.get("mask", None) saved_attn_sums = data.get("attn_sums", None) saved_density = None # GPU-only 模式没有保存 density layer_id = 0 # GPU-only 只保存 layer 0 else: # offload 模式保存的格式 print(f"[INFO] 检测到 offload 模式数据格式") Q = data["query"].to(device) K = data["key"].to(device) BSA_BLOCK_SIZE = 128 CHUNK_SIZE = 4096 STRIDE = data["stride"] THRESHOLD = data["threshold"] if isinstance(THRESHOLD, torch.Tensor): THRESHOLD = THRESHOLD[0].item() saved_mask = None saved_attn_sums = None saved_density = data.get("density", None) layer_id = data.get("layer_id", 0) batch_size, num_heads, seq_len, head_dim = Q.shape # 命令行覆盖 CHUNK_SIZE if args.chunk_size is not None: CHUNK_SIZE = args.chunk_size print(f"[INFO] 使用命令行指定的 CHUNK_SIZE={CHUNK_SIZE}") print(f"Q shape: {Q.shape}") print(f"K shape: {K.shape}") if saved_density is not None: print(f"Data layer_id: {layer_id}, saved density: {saved_density:.4f}") else: print(f"Data layer_id: {layer_id}") print(f"使用参数: STRIDE={STRIDE}, THRESHOLD={THRESHOLD}, CHUNK_SIZE={CHUNK_SIZE}, BSA_BLOCK_SIZE={BSA_BLOCK_SIZE}") print() # ============================================================ # Step 2: 使用 xattn_estimate 高层 API # ============================================================ print("=" * 60) print("Step 2: 调用 xattn_estimate (高层 API)") print("=" * 60) attn_sums_api, mask_api = xattn_estimate( Q, K, block_size=BSA_BLOCK_SIZE, stride=STRIDE, threshold=THRESHOLD, chunk_size=CHUNK_SIZE, causal=True, ) # 裁剪到有效区域 q_blocks = (seq_len + BSA_BLOCK_SIZE - 1) // BSA_BLOCK_SIZE k_blocks = (seq_len + BSA_BLOCK_SIZE - 1) // BSA_BLOCK_SIZE mask_api_valid = mask_api[:, :, :q_blocks, :k_blocks] # 计算 density (causal) causal_mask = torch.tril(torch.ones(q_blocks, k_blocks, device=device, dtype=torch.bool)) total_api = causal_mask.sum().item() * batch_size * num_heads selected_api = (mask_api_valid & causal_mask.unsqueeze(0).unsqueeze(0)).sum().item() density_api = selected_api / total_api print(f"mask_api shape (padded): {mask_api.shape}") print(f"mask_api_valid shape: {mask_api_valid.shape}") print(f"[xattn_estimate] density: {density_api:.6f} (selected={selected_api}, total={total_api})") print() # ============================================================ # Step 3: 三阶段 KV Chunking # ============================================================ print("=" * 60) print("Step 3: 三阶段 KV Chunking") print("=" * 60) print(" 1) 每个 KV chunk 计算 partial stats") print(" 2) Host 端合并 stats") print(" 3) 使用全局 stats 归一化并计算 block sums") print() # 计算 padding 参数 k_num_to_pad = ((seq_len + CHUNK_SIZE - 1) // CHUNK_SIZE) * CHUNK_SIZE - seq_len q_num_to_pad = ((seq_len + CHUNK_SIZE - 1) // CHUNK_SIZE) * CHUNK_SIZE - seq_len q_chunk_num = (seq_len + q_num_to_pad) // CHUNK_SIZE kv_chunk_num = (seq_len + k_num_to_pad) // CHUNK_SIZE k_block_num = (seq_len + k_num_to_pad) // BSA_BLOCK_SIZE q_block_num = (seq_len + q_num_to_pad) // BSA_BLOCK_SIZE reshaped_chunk_size = CHUNK_SIZE // STRIDE reshaped_block_size = BSA_BLOCK_SIZE // STRIDE k_reshaped_seq_len = (seq_len + k_num_to_pad) // STRIDE k_reshaped_num_to_pad = k_num_to_pad // STRIDE num_blocks_per_chunk = reshaped_chunk_size // reshaped_block_size kv_reshaped_chunk_size = CHUNK_SIZE // STRIDE print(f"seq_len: {seq_len}, q_chunk_num: {q_chunk_num}, kv_chunk_num: {kv_chunk_num}") print() # Padding if k_num_to_pad > 0: K_padded = torch.nn.functional.pad(K, (0, 0, 0, k_num_to_pad), value=0) else: K_padded = K if q_num_to_pad > 0: Q_padded = torch.nn.functional.pad(Q, (0, 0, 0, q_num_to_pad), value=0) else: Q_padded = Q # Softmax scale norm = 1.0 scale = 1.4426950408889634 / math.sqrt(head_dim) / STRIDE / norm simple_mask_list = [] for q_chunk_idx in range(q_chunk_num): q_start = q_chunk_idx * reshaped_chunk_size * STRIDE q_end = q_start + reshaped_chunk_size * STRIDE Q_chunk = Q_padded[:, :, q_start:q_end, :] chunk_start = (k_block_num - q_block_num) * reshaped_block_size + q_chunk_idx * reshaped_chunk_size chunk_end = chunk_start + reshaped_chunk_size # 阶段 1: 每个 KV chunk 计算 partial stats 和 raw scores m_chunks = [] l_chunks = [] attn_weights_chunks = [] for kv_chunk_idx in range(kv_chunk_num): kv_start = kv_chunk_idx * CHUNK_SIZE kv_end = kv_start + CHUNK_SIZE K_chunk = K_padded[:, :, kv_start:kv_end, :] # KV offset in reshaped space kv_offset_reshaped = kv_chunk_idx * kv_reshaped_chunk_size # 计算 raw attention scores attn_weights_kv = flat_group_gemm_fuse_reshape( Q_chunk, K_chunk, STRIDE, chunk_start=chunk_start, chunk_end=chunk_end, is_causal=False, # K 不完整，不能在这里用 causal ) attn_weights_chunks.append(attn_weights_kv) # 计算 partial stats (带 causal mask) m_partial, l_partial = softmax_compute_partial_stats( attn_weights_kv, reshaped_block_size, min(4096, reshaped_block_size), scale, chunk_start=chunk_start, kv_offset=kv_offset_reshaped, is_causal=True, ) m_chunks.append(m_partial) l_chunks.append(l_partial) # 阶段 2: Host 端合并 stats m_global, l_global = merge_softmax_stats(m_chunks, l_chunks) # 阶段 3: 使用全局 stats 归一化并计算 block sums attn_sum_per_kv = [] for kv_chunk_idx, attn_weights_kv in enumerate(attn_weights_chunks): kv_offset_reshaped = kv_chunk_idx * kv_reshaped_chunk_size attn_sum_kv = softmax_normalize_and_block_sum( attn_weights_kv, m_global, l_global, reshaped_block_size, min(4096, reshaped_block_size), chunk_start=chunk_start, real_q_len=k_reshaped_seq_len - k_reshaped_num_to_pad, scale=scale, kv_offset=kv_offset_reshaped, is_causal=True, ) attn_sum_per_kv.append(attn_sum_kv) # 拼接各 KV chunk 的 block sums attn_sum_concat = torch.cat(attn_sum_per_kv, dim=-1) # 选择 blocks simple_mask = find_blocks_chunked( attn_sum_concat, current_index=k_block_num - q_block_num + q_chunk_idx * num_blocks_per_chunk, threshold=THRESHOLD, num_to_choose=None, decoding=False, mode="prefill", causal=True, ) simple_mask_list.append(simple_mask) print(f" Q chunk {q_chunk_idx}: merged {kv_chunk_num} KV chunks, attn_sum shape={attn_sum_concat.shape}") mask_kv_chunking = torch.cat(simple_mask_list, dim=2) # 应用与 xattn_estimate 相同的 causal mask 后处理 (xattn.py 第 1300-1306 行) mask_kv_chunking[:, :, -q_block_num:, -q_block_num:] = torch.where( torch.tril(torch.ones(q_block_num, q_block_num, dtype=bool, device=device), diagonal=0), mask_kv_chunking[:, :, -q_block_num:, -q_block_num:], False, ) mask_kv_chunking_valid = mask_kv_chunking[:, :, :q_blocks, :k_blocks] selected_kv = (mask_kv_chunking_valid & causal_mask.unsqueeze(0).unsqueeze(0)).sum().item() density_kv = selected_kv / total_api print() print(f"[KV chunking] density: {density_kv:.6f} (selected={selected_kv}, total={total_api})") print() # ============================================================ # Step 4: 对比结果 # ============================================================ print("=" * 60) print("Step 4: 对比结果") print("=" * 60) print() mask_total = mask_api_valid.numel() mask_diff = (mask_api_valid != mask_kv_chunking_valid).sum().item() print("| 方法 | density | 与 API 差异 | Mask 差异 |") print("|------|---------|-------------|-----------|") print(f"| xattn_estimate API | {density_api:.6f} | - | - |") print(f"| KV chunking | {density_kv:.6f} | {abs(density_api - density_kv):.6f} | {100*mask_diff/mask_total:.4f}% |") print() passed = abs(density_api - density_kv) < 1e-6 and mask_diff / mask_total < 0.001 # ============================================================ # Step 5: 与 GPU-only 保存的数据对比 (如果有) # ============================================================ if saved_mask is not None or saved_attn_sums is not None: print("=" * 60) print("Step 5: 与 GPU-only 保存的数据对比") print("=" * 60) print() if saved_mask is not None: saved_mask_gpu = saved_mask.to(device) # 比较 mask mask_saved_diff = (mask_api_valid != saved_mask_gpu).sum().item() mask_saved_total = saved_mask_gpu.numel() print(f"| xattn_estimate vs GPU-only saved mask | 差异 blocks: {mask_saved_diff} / {mask_saved_total} ({100*mask_saved_diff/mask_saved_total:.4f}%) |") if mask_saved_diff == 0: print("✅ mask 与 GPU-only 保存完全一致") else: print("❌ mask 与 GPU-only 保存存在差异") passed = False if saved_attn_sums is not None: saved_attn_sums_gpu = saved_attn_sums.to(device) # 需要从 xattn_estimate 获取 attn_sums # 重新调用一次获取 attn_sums attn_sums_check, _ = xattn_estimate( Q, K, block_size=BSA_BLOCK_SIZE, stride=STRIDE, threshold=THRESHOLD, chunk_size=CHUNK_SIZE, causal=True, ) attn_sums_check_valid = attn_sums_check[:, :, :q_blocks, :k_blocks] max_diff = (attn_sums_check_valid - saved_attn_sums_gpu).abs().max().item() mean_diff = (attn_sums_check_valid - saved_attn_sums_gpu).abs().mean().item() print(f"| xattn_estimate vs GPU-only saved attn_sums | max diff: {max_diff:.6e}, mean diff: {mean_diff:.6e} |") if max_diff < 1e-5: print("✅ attn_sums 与 GPU-only 保存一致") else: print("❌ attn_sums 与 GPU-only 保存存在差异") passed = False print() if passed: print("test_xattn_estimate_alignment: PASSED") else: print("test_xattn_estimate_alignment: FAILED")