Flash Decoding
长上下文推理加速

# 标准实现：需要物化 N×N 的中间矩阵
S = Q @ K.T / sqrt(d)          # HBM写入: N×N 的 float16 = 2N² 字节
P = softmax(S)                  # HBM读取 S，写入 P：再 2N² 字节
O = P @ V                       # HBM读取 P：再 2N² 字节
# 总 HBM 访问量：O(N²d) → 对长序列极其昂贵

@triton.jit
def flash_decode_kernel1(
    Q_ptr, K_ptr, V_ptr,
    Out_ptr, LSE_ptr,    # LSE = log-sum-exp
    seq_len, chunk_size,
    stride_qh, stride_kh, stride_vh,
    BLOCK_D: tl.constexpr,
):
    # 当前 thread block 负责的 split 索引
    split_id = tl.program_id(0)
    head_id  = tl.program_id(1)
    
    # 加载 Q (1 × d)
    q = tl.load(Q_ptr + head_id * stride_qh + tl.arange(0, BLOCK_D))
    
    # 初始化局部状态
    m_i = tl.full((1,), -1e9, dtype=tl.float32)  # running max
    l_i = tl.zeros((1,), dtype=tl.float32)        # running sum
    acc = tl.zeros((BLOCK_D,), dtype=tl.float32)    # running output
    
    # 在当前 split 内部做 tiling（FlashAttention 风格）
    start = split_id * chunk_size
    for i in range(0, chunk_size, BLOCK_C):
        # 加载 K block, V block
        k = tl.load(K_ptr + (start + i) * stride_kh + tl.arange(0, BLOCK_D))
        v = tl.load(V_ptr + (start + i) * stride_vh + tl.arange(0, BLOCK_D))
        
        # 计算 QK score
        s = tl.dot(q[None,:], k[:,None]) * scale  # (1,)
        
        # Online softmax 更新
        m_new = tl.maximum(m_i, s)
        l_i = l_i * tl.exp(m_i - m_new) + tl.exp(s - m_new)
        acc  = acc * tl.exp(m_i - m_new) + tl.exp(s - m_new) * v
        m_i  = m_new
    
    # 归一化并写回
    acc = acc / l_i
    lse = tl.log(l_i) + m_i   # log-sum-exp
    tl.store(Out_ptr + ..., acc)
    tl.store(LSE_ptr + ..., lse)

@triton.jit
def flash_decode_kernel2(
    Out_splits_ptr, LSE_splits_ptr,   # Kernel 1 的输出
    Final_out_ptr,
    num_splits,
    BLOCK_D: tl.constexpr,
):
    head_id = tl.program_id(0)
    d_ids   = tl.arange(0, BLOCK_D)
    
    # Step 1: 读取所有 split 的 LSE 和 out
    lse_all = tl.load(LSE_splits_ptr + head_id * num_splits 
                      + tl.arange(0, num_splits))  # (S,)
    
    # Step 2: 计算全局 log-sum-exp
    lse_global = tl.reduce(lse_all, 0, tl.math.logsumexp)
    
    # Step 3: 加权合并
    acc = tl.zeros((BLOCK_D,), dtype=tl.float32)
    for s in range(num_splits):
        lse_s = tl.load(LSE_splits_ptr + head_id * num_splits + s)
        out_s = tl.load(Out_splits_ptr + (head_id * num_splits + s) * BLOCK_D + d_ids)
        scale = tl.exp(lse_s - lse_global)
        acc  += scale * out_s
    
    # Step 4: 写回
    tl.store(Final_out_ptr + head_id * BLOCK_D + d_ids, acc)

# flash_attn/flash_attn_interface.py
def flash_attn_with_kvcache(
    q,          # (batch, seqlen_q, nheads, headdim)
    k_cache,    # (batch, seqlen_k, nheads_k, headdim)
    v_cache,    # (batch, seqlen_k, nheads_k, headdim)
    cache_seqlens=None,   # 每个 batch 的实际 KV 长度
    softmax_scale=None,
    causal=False,
    num_splits=0,        # 0 = 自动选择；>0 = 手动指定 split 数
):
    """
    Flash Decoding 接口。
    当 seqlen_q == 1 时（典型 decoding 场景）自动启用 Flash Decoding。
    支持 MQA/GQA（nheads_k < nheads）。
    """
    if num_splits == 0:
        # 启发式：根据 seq_len 和 SM 数量自动选择 split 数
        num_splits = _get_num_splits(seqlen_k, nheads, batch)
    
    # Kernel 1：并行计算各 split 的局部 attention
    out_splits, lse_splits = _flash_decode_split(
        q, k_cache, v_cache, num_splits, softmax_scale)
    
    # Kernel 2：归约合并
    out = _flash_decode_reduce(out_splits, lse_splits)
    
    return out

import triton
import triton.language as tl
import torch

@triton.jit
def _flash_decode_fwd_kernel(
    Q, K, V,           # [B, H, 1, D], [B, H_k, N, D], [B, H_k, N, D]
    Out, LSE,          # [B, H, S, D], [B, H, S]  (S = num_splits)
    stride_qb, stride_qh, stride_qd,
    stride_kb, stride_kh, stride_kn, stride_kd,
    stride_ob, stride_oh, stride_os, stride_od,
    stride_lb, stride_lh, stride_ls,
    N,                 # 总序列长度
    num_splits,
    scale,
    BLOCK_D: tl.constexpr,   # head dim，必须是 power of 2
    BLOCK_N: tl.constexpr,   # 每次处理的 K 行数（tile size）
    GROUP_H: tl.constexpr,   # GQA: Q heads per KV head
):
    # ── 索引计算 ─────────────────────────────────────────
    batch_id = tl.program_id(0)
    head_id  = tl.program_id(1)
    split_id = tl.program_id(2)
    
    kv_head_id = head_id // GROUP_H   # GQA: 映射到 KV head
    
    chunk_size = tl.cdiv(N, num_splits)
    start = split_id * chunk_size
    end   = tl.minimum(start + chunk_size, N)
    
    # ── 加载 Q ──────────────────────────────────────────
    q = tl.load(
        Q + batch_id * stride_qb + head_id * stride_qh 
          + tl.arange(0, BLOCK_D),
        mask=tl.arange(0, BLOCK_D) < BLOCK_D, other=0.0
    ).to(tl.float32)
    
    # ── 初始化 Online Softmax 状态 ──────────────────────
    m_i = tl.full((1,), float("-inf"), dtype=tl.float32)
    l_i = tl.zeros((1,), dtype=tl.float32)
    acc = tl.zeros((BLOCK_D,), dtype=tl.float32)
    
    # ── 在当前 split 内部 tile 循环 ──────────────────────
    kv_base_k = (batch_id * stride_kb + kv_head_id * stride_kh 
                 + start * stride_kn)
    kv_base_v = (batch_id * stride_kb + kv_head_id * stride_kh 
                 + start * stride_kn)
    
    for block_start in range(start, end, BLOCK_N):
        block_end = tl.minimum(block_start + BLOCK_N, end)
        valid = block_end - block_start
        
        n_ids = tl.arange(0, BLOCK_N)
        mask_n = n_ids < valid
        
        # 加载 K block: [BLOCK_N, BLOCK_D]
        k = tl.load(
            K + kv_base_k + n_ids[:, None] * stride_kn 
              + tl.arange(0, BLOCK_D)[None, :],
            mask=mask_n[:, None], other=0.0
        ).to(tl.float32)
        
        # 加载 V block: [BLOCK_N, BLOCK_D]
        v = tl.load(
            V + kv_base_v + n_ids[:, None] * stride_kn 
              + tl.arange(0, BLOCK_D)[None, :],
            mask=mask_n[:, None], other=0.0
        ).to(tl.float32)
        
        # QK^T：scores [BLOCK_N]
        scores = tl.sum(q[None, :] * k, axis=1) * scale
        scores = tl.where(mask_n, scores, float("-inf"))
        
        # Online Softmax 更新
        m_new = tl.maximum(m_i, tl.max(scores, axis=0))
        p = tl.exp(scores - m_new)           # [BLOCK_N]
        l_new = l_i * tl.exp(m_i - m_new) + tl.sum(p, axis=0)
        acc = acc * tl.exp(m_i - m_new) + tl.sum(p[:, None] * v, axis=0)
        m_i, l_i = m_new, l_new
        
        kv_base_k += BLOCK_N * stride_kn
        kv_base_v += BLOCK_N * stride_kn
    
    # ── 写回局部结果 ─────────────────────────────────────
    acc = acc / l_i   # 归一化
    lse = tl.log(l_i) + m_i   # log-sum-exp
    
    out_off = (batch_id * stride_ob + head_id * stride_oh 
               + split_id * stride_os + tl.arange(0, BLOCK_D))
    tl.store(Out + out_off, acc.to(tl.float16))
    
    lse_off = batch_id * stride_lb + head_id * stride_lh + split_id * stride_ls
    tl.store(LSE + lse_off, lse)


@triton.jit
def _flash_decode_reduce_kernel(
    Out_splits, LSE_splits,   # [B, H, S, D], [B, H, S]
    Final_out,                # [B, H, D]
    num_splits,
    BLOCK_D: tl.constexpr,
):
    batch_id = tl.program_id(0)
    head_id  = tl.program_id(1)
    
    # 读取所有 split 的 LSE
    lse_off = batch_id * (num_splits * ...) + head_id * num_splits
    lse_vals = tl.load(LSE_splits + lse_off + tl.arange(0, num_splits))
    
    # 全局 log-sum-exp：数值稳定实现
    lse_max = tl.max(lse_vals, axis=0)
    lse_global = lse_max + tl.log(tl.sum(tl.exp(lse_vals - lse_max), axis=0))
    
    # 加权合并所有 split 的 output
    acc = tl.zeros((BLOCK_D,), dtype=tl.float32)
    for s in range(num_splits):
        w = tl.exp(tl.load(LSE_splits + lse_off + s) - lse_global)
        out_s = tl.load(Out_splits + ... + s * BLOCK_D + tl.arange(0, BLOCK_D))
        acc += w * out_s.to(tl.float32)
    
    tl.store(Final_out + batch_id * ... + head_id * BLOCK_D 
             + tl.arange(0, BLOCK_D), acc.to(tl.float16))

def flash_decode(q, k_cache, v_cache, num_splits=None):
    """
    q:       [batch, 1, nheads, headdim]       ← seqlen=1 for decoding
    k_cache: [batch, seqlen, nheads_k, headdim] ← KV Cache
    v_cache: [batch, seqlen, nheads_k, headdim]
    """
    B, _, H, D = q.shape
    _, N, H_k, _ = k_cache.shape
    GROUP_H = H // H_k   # GQA group size
    
    if num_splits is None:
        # 启发式：让总 Thread Block 数 ≈ SM 数的 4 倍（充分利用）
        num_sm = torch.cuda.get_device_properties(0).multi_processor_count
        num_splits = max(1, num_sm * 4 // (B * H))
        num_splits = min(num_splits, N // 64)   # 每个 split 至少 64 个 token
    
    scale = D ** -0.5
    
    # 临时缓冲区
    out_splits = torch.empty(B, H, num_splits, D, dtype=torch.float16, 
                             device=q.device)
    lse_splits  = torch.empty(B, H, num_splits, dtype=torch.float32, 
                              device=q.device)
    
    # Kernel 1 启动
    grid1 = (B, H, num_splits)
    _flash_decode_fwd_kernel[grid1](
        q, k_cache, v_cache, out_splits, lse_splits,
        *q.stride(), *k_cache.stride(),
        *out_splits.stride(), *lse_splits.stride(),
        N, num_splits, scale,
        BLOCK_D=D, BLOCK_N=64, GROUP_H=GROUP_H,
    )
    
    # Kernel 2 启动（归约）
    final_out = torch.empty(B, H, D, dtype=torch.float16, device=q.device)
    grid2 = (B, H)
    _flash_decode_reduce_kernel[grid2](
        out_splits, lse_splits, final_out, num_splits, BLOCK_D=D)
    
    return final_out.unsqueeze(1)  # [B, 1, H, D]

from flash_attn import flash_attn_with_kvcache

class LlamaAttention:
    def forward(self, x, kv_cache, position_ids):
        B, seq_q, _ = x.shape
        
        q = self.q_proj(x).view(B, seq_q, self.n_heads, self.head_dim)
        k = self.k_proj(x).view(B, seq_q, self.n_kv_heads, self.head_dim)
        v = self.v_proj(x).view(B, seq_q, self.n_kv_heads, self.head_dim)
        
        # 更新 KV Cache
        kv_cache.k[..., :seq_q, :] = k
        kv_cache.v[..., :seq_q, :] = v
        
        # Flash Decoding (当 seq_q == 1 时自动启用)
        out = flash_attn_with_kvcache(
            q,              # [B, 1, H, D]
            kv_cache.k,     # [B, N, H_k, D]
            kv_cache.v,     # [B, N, H_k, D]
            cache_seqlens=kv_cache.seq_lens,  # 每条序列的实际长度
            softmax_scale=self.head_dim**-0.5,
            causal=True,
        )
        # out: [B, 1, H, D]
        
        out = out.view(B, seq_q, self.n_heads * self.head_dim)
        return self.o_proj(out)

import xformers.ops as xops

# xFormers 会自动 dispatch：
# - seqlen_q > 1 → FlashAttention（训练/prefill）
# - seqlen_q == 1 → Flash Decoding（decoding）
out = xops.memory_efficient_attention(
    q, k, v,
    attn_bias=None,
    scale=head_dim**-0.5,
)