Update

Author: Hanyuezhuohua

analyze_compression_diff.py

@@ -0,0 +1,256 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+分析 compression branch 的 naive 和 triton 实现差异
+重点关注为什么 backward dq 的值不同
+"""
+
+import os
+import torch
+import triton
+
+os.environ['TRITON_F32_DEFAULT'] = 'ieee'
+
+from native_sparse_attention.ops.parallel import parallel_nsa_compression
+from fla.ops.utils import mean_pooling
+
+
+def analyze_forward_differences():
+    """分析 forward 的实现差异"""
+    print("="*80)
+    print("分析 Forward 实现差异")
+    print("="*80)
+
+    B, T, H, HQ, D = 1, 256, 4, 64, 64
+    block_size = 32
+    dtype = torch.bfloat16
+    scale = 0.1
+
+    torch.manual_seed(42)
+
+    # 创建输入
+    q = torch.randn((B, T, HQ, D), dtype=dtype, device='cuda')
+    k = torch.randn((B, T, H, D), dtype=dtype, device='cuda')
+    v = torch.randn((B, T, H, D), dtype=dtype, device='cuda')
+
+    # Compression
+    k_cmp = mean_pooling(k, block_size, None)
+    v_cmp = mean_pooling(v, block_size, None)
+
+    print(f"\n1. 输入形状:")
+    print(f"   q: {q.shape}")
+    print(f"   k: {k.shape} -> k_cmp: {k_cmp.shape}")
+    print(f"   v: {v.shape} -> v_cmp: {v_cmp.shape}")
+
+    # Triton 实现
+    print(f"\n2. Triton 实现 (parallel_nsa_compression):")
+    print(f"   - 使用优化的 Triton kernel")
+    print(f"   - Forward kernel: parallel_nsa_compression_fwd_kernel")
+    print(f"   - 特点:")
+    print(f"     * 在线计算 softmax (online softmax algorithm)")
+    print(f"     * 使用 log-sum-exp trick 保持数值稳定")
+    print(f"     * 分块处理，减少内存访问")
+
+    o_tri, lse_tri = parallel_nsa_compression(
+        q=q, k=k_cmp, v=v_cmp,
+        block_size=block_size, scale=scale, offsets=None
+    )
+
+    print(f"   - 输出: o shape={o_tri.shape}, lse shape={lse_tri.shape}")
+    print(f"   - o 的范围: [{o_tri.min().item():.4f}, {o_tri.max().item():.4f}]")
+    print(f"   - lse 的范围: [{lse_tri.min().item():.4f}, {lse_tri.max().item():.4f}]")
+
+    # Naive 实现 (手工实现标准attention)
+    print(f"\n3. Naive 实现 (手工标准 attention):")
+    print(f"   - 使用 PyTorch 原生操作")
+    print(f"   - 直接计算: softmax(Q @ K^T) @ V")
+
+    # 手工实现 compression attention
+    G = HQ // H
+    k_cmp_expanded = k_cmp.repeat_interleave(G, dim=2)  # [B, C, HQ, D]
+    v_cmp_expanded = v_cmp.repeat_interleave(G, dim=2)
+
+    q_float = q.float()
+    k_cmp_float = k_cmp_expanded.float()
+    v_cmp_float = v_cmp_expanded.float()
+
+    # 重排维度用于 attention: [B, HQ, T, D]
+    q_attn = q_float.transpose(1, 2)  # [B, HQ, T, D]
+    k_attn = k_cmp_float.transpose(1, 2)  # [B, HQ, C, D]
+    v_attn = v_cmp_float.transpose(1, 2)  # [B, HQ, C, D]
+
+    # Attention scores
+    C = k_cmp.shape[1]
+    attn_scores = torch.matmul(q_attn * scale, k_attn.transpose(-2, -1))  # [B, HQ, T, C]
+
+    # Causal mask
+    causal_mask = ((torch.arange(T, device='cuda') - block_size + 1)[:, None] // block_size
+                   < torch.arange(C, device='cuda')[None, :])
+    empty_mask = causal_mask.all(-1, keepdim=True)
+
+    attn_scores = attn_scores.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), float('-inf'))
+    attn_probs = torch.softmax(attn_scores, dim=-1)
+    attn_probs = attn_probs.masked_fill(empty_mask.unsqueeze(0).unsqueeze(0), 0.0)
+
+    o_naive = torch.matmul(attn_probs, v_attn)  # [B, HQ, T, D]
+    o_naive = o_naive.transpose(1, 2).to(dtype)  # [B, T, HQ, D]
+
+    # LSE for naive
+    lse_naive = torch.logsumexp(attn_scores, dim=-1).to(torch.float32)  # [B, HQ, T]
+    lse_naive = lse_naive.transpose(1, 2)  # [B, T, HQ]
+
+    print(f"   - 输出: o shape={o_naive.shape}, lse shape={lse_naive.shape}")
+    print(f"   - o 的范围: [{o_naive.min().item():.4f}, {o_naive.max().item():.4f}]")
+    print(f"   - lse 的范围: [{lse_naive.min().item():.4f}, {lse_naive.max().item():.4f}]")
+
+    # 比较
+    print(f"\n4. Forward 输出对比:")
+    o_diff = (o_tri - o_naive).abs()
+    lse_diff = (lse_tri - lse_naive).abs()
+
+    print(f"   - o 最大差异: {o_diff.max().item():.6f}")
+    print(f"   - o 平均差异: {o_diff.mean().item():.6f}")
+    print(f"   - o 相对误差: {(o_diff.mean() / o_naive.abs().mean()).item():.6f}")
+
+    print(f"   - lse 最大差异: {lse_diff.max().item():.6f}")
+    print(f"   - lse 平均差异: {lse_diff.mean().item():.6f}")
+
+    return q, k_cmp, v_cmp, o_tri, o_naive, lse_tri, lse_naive
+
+
+def analyze_backward_differences(q, k_cmp, v_cmp, o_tri, o_naive, lse_tri, lse_naive):
+    """分析 backward 的实现差异"""
+    print("\n" + "="*80)
+    print("分析 Backward 实现差异")
+    print("="*80)
+
+    B, T, HQ, D = q.shape
+    H = k_cmp.shape[2]
+    G = HQ // H
+    block_size = 32
+    scale = 0.1
+
+    # 创建梯度输出
+    do = torch.randn_like(o_tri)
+
+    print(f"\n1. Backward 输入:")
+    print(f"   - do shape: {do.shape}")
+    print(f"   - do 的范围: [{do.min().item():.4f}, {do.max().item():.4f}]")
+
+    # Triton backward
+    print(f"\n2. Triton Backward:")
+    print(f"   - 使用 parallel_nsa_compression_bwd_kernel_dq")
+    print(f"   - 关键步骤:")
+    print(f"     * delta = sum(do * o, dim=-1)  # [B, T, HQ]")
+    print(f"     * p = exp(s - lse)  # 重新计算 attention probs")
+    print(f"     * ds = p * (dp - delta)")
+    print(f"     * dq = ds @ K^T * scale")
+    print(f"   - 特点: 使用 FlashAttention-style backward")
+
+    q_tri = q.clone().requires_grad_(True)
+    k_cmp_tri = k_cmp.detach()
+    v_cmp_tri = v_cmp.detach()
+
+    o_tri_new, _ = parallel_nsa_compression(
+        q=q_tri, k=k_cmp_tri, v=v_cmp_tri,
+        block_size=block_size, scale=scale, offsets=None
+    )
+    o_tri_new.backward(do)
+    dq_tri = q_tri.grad.clone()
+
+    print(f"   - dq shape: {dq_tri.shape}")
+    print(f"   - dq 的范围: [{dq_tri.min().item():.4f}, {dq_tri.max().item():.4f}]")
+    print(f"   - dq norm: {dq_tri.norm().item():.6f}")
+
+    # Naive backward
+    print(f"\n3. Naive Backward (PyTorch autograd):")
+    print(f"   - 使用 PyTorch 自动微分")
+    print(f"   - 标准 attention backward公式")
+
+    q_naive = q.clone().requires_grad_(True)
+    k_cmp_naive = k_cmp.clone().detach()
+    v_cmp_naive = v_cmp.clone().detach()
+
+    # 重新计算 forward
+    k_cmp_expanded = k_cmp_naive.repeat_interleave(G, dim=2)
+    v_cmp_expanded = v_cmp_naive.repeat_interleave(G, dim=2)
+
+    q_float = q_naive.float()
+    k_float = k_cmp_expanded.float()
+    v_float = v_cmp_expanded.float()
+
+    q_attn = q_float.transpose(1, 2)
+    k_attn = k_float.transpose(1, 2)
+    v_attn = v_float.transpose(1, 2)
+
+    C = k_cmp_naive.shape[1]
+    attn_scores = torch.matmul(q_attn * scale, k_attn.transpose(-2, -1))
+
+    causal_mask = ((torch.arange(T, device='cuda') - block_size + 1)[:, None] // block_size
+                   < torch.arange(C, device='cuda')[None, :])
+    empty_mask = causal_mask.all(-1, keepdim=True)
+
+    attn_scores = attn_scores.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), float('-inf'))
+    attn_probs = torch.softmax(attn_scores, dim=-1)
+    attn_probs = attn_probs.masked_fill(empty_mask.unsqueeze(0).unsqueeze(0), 0.0)
+
+    o_naive_new = torch.matmul(attn_probs, v_attn)
+    o_naive_new = o_naive_new.transpose(1, 2).to(q_naive.dtype)
+
+    o_naive_new.backward(do)
+    dq_naive = q_naive.grad.clone()
+
+    print(f"   - dq shape: {dq_naive.shape}")
+    print(f"   - dq 的范围: [{dq_naive.min().item():.4f}, {dq_naive.max().item():.4f}]")
+    print(f"   - dq norm: {dq_naive.norm().item():.6f}")
+
+    # 比较
+    print(f"\n4. Backward dq 对比:")
+    dq_diff = (dq_tri - dq_naive).abs()
+
+    print(f"   - 最大差异: {dq_diff.max().item():.6f}")
+    print(f"   - 平均差异: {dq_diff.mean().item():.6f}")
+    print(f"   - 相对误差: {(dq_diff.norm() / dq_naive.norm()).item():.6f}")
+    print(f"   - Norm 比率 (tri/naive): {(dq_tri.norm() / dq_naive.norm()).item():.6f}")
+
+    # 详细分析差异来源
+    print(f"\n5. 差异来源分析:")
+    print(f"   可能的原因:")
+    print(f"   a) 数值精度: Triton kernel 使用 bfloat16, naive 使用 float32")
+    print(f"   b) 计算顺序: Triton 的在线算法 vs PyTorch 的批量计算")
+    print(f"   c) Softmax 实现: FlashAttention-style vs 标准 softmax")
+    print(f"   d) 舍入误差累积: 分块计算可能导致不同的舍入模式")
+
+    # 检查 attention 模式是否一致
+    print(f"\n6. 检查 attention 模式:")
+    with torch.no_grad():
+        q_test = q.float()
+        k_test = k_cmp_expanded.float()
+
+        q_test_attn = q_test.transpose(1, 2)
+        k_test_attn = k_test.transpose(1, 2)
+
+        scores_test = torch.matmul(q_test_attn * scale, k_test_attn.transpose(-2, -1))
+        scores_test = scores_test.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), float('-inf'))
+        probs_test = torch.softmax(scores_test, dim=-1)
+
+        # 检查几个样本位置的 attention 分布
+        for t in [50, 100, 150, 200]:
+            print(f"   位置 t={t}:")
+            prob_sample = probs_test[0, 0, t, :t//block_size]
+            print(f"     - Top-5 attention weights: {prob_sample.topk(5)[0].cpu().numpy()}")
+            print(f"     - Entropy: {-(prob_sample * torch.log(prob_sample + 1e-10)).sum().item():.4f}")
+
+
+if __name__ == '__main__':
+    q, k_cmp, v_cmp, o_tri, o_naive, lse_tri, lse_naive = analyze_forward_differences()
+    analyze_backward_differences(q, k_cmp, v_cmp, o_tri, o_naive, lse_tri, lse_naive)
+
+    print("\n" + "="*80)
+    print("结论:")
+    print("="*80)
+    print("compression branch 的 backward dq 存在显著差异，主要原因可能是:")
+    print("1. Triton kernel 的在线 softmax 算法与标准实现的数值差异")
+    print("2. 混合精度计算带来的累积误差")
+    print("3. 分块计算的舍入误差传播")
+    print("\n建议检查 Triton backward kernel 的实现是否正确匹配 forward 的计算逻辑")