1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101

@triton.jit
def Decode_Paged_GQAAttention_Kernel(
        q                         ,  # (B, n_heads, 1, head_dim)         query，decode 每步只有 1 个 token
        k_cache                   ,  # (num_blocks, n_kv_heads, block_size, head_dim)  全局 K cache
        v_cache                   ,  # (num_blocks, n_kv_heads, block_size, head_dim)  全局 V cache
        block_table               ,  # (B, MAX_BLOCKS_PER_SEQ)  int32，每个请求的物理块 id，-1 表示未分配
        context_lens              ,  # (B,)             int32，每个请求当前的有效 token 数
        scale              ,         # 1.0 / sqrt(head_dim)
        out,
        
        # Meta-parameters
        # 元参数
        MAX_BLOCKS_PER_SEQ: tl.constexpr,  # 启动时固定，不是运行时变量
        BLOCK_SIZE: tl.constexpr,  #
        HEAD_DIM: tl.constexpr,  #
        N_KV_HEAD: tl.constexpr,
        N_HEAD: tl.constexpr,
    ):               # (B, n_heads, 1, head_dim)
    # grid = (B, n_heads)
    # 每个 program 处理一个 (batch, head) 对
    # program_id[0] = batch_idx
    # program_id[1] = head_idx
    
    pid_batch = tl.program_id(axis=0)
    pid_head = tl.program_id(axis=1)
    kv_head_idx = (pid_head // (N_HEAD // N_KV_HEAD))

    m = float('-inf')
    l = float(0)
    o = tl.zeros((HEAD_DIM, ), dtype=tl.float32)
    
    
    q_ptrs = q + pid_batch * (HEAD_DIM * N_HEAD) + pid_head * (HEAD_DIM) + tl.arange(0, HEAD_DIM)
    q_vec = tl.reshape(tl.load(q_ptrs), (HEAD_DIM, 1))

    context_len = tl.load(context_lens + pid_batch)
    if context_len == 0:
        return
    max_block_index = tl.cdiv(context_len, BLOCK_SIZE)  # 向上取整
    offs_kv = tl.arange(0, BLOCK_SIZE * HEAD_DIM)
    
    for block_idx in range(0, max_block_index):
        
        physical_idx = tl.load(block_table + pid_batch * MAX_BLOCKS_PER_SEQ + block_idx)
        physical_idx = tl.maximum(physical_idx, 0).to(tl.int64)
        base = (physical_idx * N_KV_HEAD * BLOCK_SIZE * HEAD_DIM + kv_head_idx* BLOCK_SIZE * HEAD_DIM)

        # 加载时 mask 掉超出 context_len 的 token
        token_start = block_idx * BLOCK_SIZE
        valid_in_block = tl.minimum(BLOCK_SIZE, context_len - token_start)
        kv_token_mask = tl.arange(0, BLOCK_SIZE * HEAD_DIM) < valid_in_block * HEAD_DIM

        k_ptrs = k_cache + base + offs_kv
        # k_block: (block_size, head_dim)
        k_block = tl.load(k_ptrs, mask = kv_token_mask, other=0.0)
        k_block = tl.reshape(k_block, (BLOCK_SIZE, HEAD_DIM))
        v_ptrs = v_cache + base + offs_kv
        v_block = tl.load(v_ptrs, mask = kv_token_mask, other=0.0)
        v_block = tl.reshape(v_block, (BLOCK_SIZE, HEAD_DIM))

        # mask 最后一个 block 的无效 token
        valid = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) < context_len
        q_row = tl.reshape(q_vec, (1, HEAD_DIM))           # (1, HEAD_DIM)
        s = tl.sum(k_block * q_row, axis=1)                 # (BLOCK_SIZE,)
        s = s.to(tl.float32) * scale
        s = tl.where(valid, s, float('-inf'))

        m_new = tl.maximum(m, tl.max(s))
        alpha = tl.exp(m - m_new)           # 旧的缩放因子
        p = tl.exp(s - m_new)          # 当前 block 的权重

        l = l * alpha + tl.sum(p)
        p_col = tl.reshape(p, (BLOCK_SIZE, 1))             # (BLOCK_SIZE, 1)
        o = o * alpha + tl.sum(p_col * v_block, axis=0)    # (HEAD_DIM,)
        m = m_new

    o = o / l  # (HEAD_DIM,)
    o_casted = tl.cast(o, out.dtype.element_ty)
    
    tl.store(out + pid_batch * (HEAD_DIM * N_HEAD) + pid_head * (HEAD_DIM) + tl.arange(0, HEAD_DIM), o_casted)

@torch.compiler.disable
def paged_decode_attention(q, k_cache, v_cache, block_table, context_lens,
                           MAX_BLOCKS_PER_SEQ, BLOCK_SIZE=16):
    B, N_HEAD, _, HEAD_DIM = q.shape
    N_KV_HEAD = k_cache.shape[1]
    scale = 1.0 / (HEAD_DIM ** 0.5)

    out = torch.empty(B, N_HEAD, 1, HEAD_DIM, device=q.device, dtype=q.dtype)

    grid = (B, N_HEAD)
    Decode_Paged_GQAAttention_Kernel[grid](
        q, k_cache, v_cache, block_table, context_lens,
        scale, out,
        MAX_BLOCKS_PER_SEQ=MAX_BLOCKS_PER_SEQ,
        BLOCK_SIZE=BLOCK_SIZE,
        HEAD_DIM=HEAD_DIM,
        N_KV_HEAD=N_KV_HEAD,
        N_HEAD=N_HEAD,
    )
    return out

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

@triton.jit
def _flash_decode_partial_kernel(
    q,                  # (B, N_HEAD, HEAD_DIM)  (decode: the seq=1 dim collapsed)
    k_cache,            # (num_blocks, N_KV_HEAD, BLOCK_SIZE, HEAD_DIM)
    v_cache,
    block_table,        # (B, MAX_BLOCKS_PER_SEQ) int32
    context_lens,       # (B,) int32
    scale,
    partial_acc,        # (B, N_HEAD, NUM_SPLITS, HEAD_DIM) fp32
    partial_m,          # (B, N_HEAD, NUM_SPLITS) fp32
    partial_l,          # (B, N_HEAD, NUM_SPLITS) fp32
    MAX_BLOCKS_PER_SEQ: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
    HEAD_DIM: tl.constexpr,
    N_KV_HEAD: tl.constexpr,
    N_HEAD: tl.constexpr,
    GROUP_SIZE: tl.constexpr,
    PADDED_GROUP: tl.constexpr,
    NUM_SPLITS: tl.constexpr,
):
    pid_b = tl.program_id(0)
    pid_kv = tl.program_id(1)
    pid_s = tl.program_id(2)

    ctx = tl.load(context_lens + pid_b)
    q_head0 = pid_kv * GROUP_SIZE

    row = tl.arange(0, PADDED_GROUP)          # (PG,)
    row_mask = row < GROUP_SIZE
    d = tl.arange(0, HEAD_DIM)                 # (HD,)
    qheads = q_head0 + row                     # (PG,) absolute q-head ids

    # partial output pointers for this (b, group, split)
    pm_ptr = partial_m + pid_b * (N_HEAD * NUM_SPLITS) + qheads * NUM_SPLITS + pid_s
    pl_ptr = partial_l + pid_b * (N_HEAD * NUM_SPLITS) + qheads * NUM_SPLITS + pid_s
    pacc_ptr = (partial_acc
                + pid_b * (N_HEAD * NUM_SPLITS * HEAD_DIM)
                + qheads[:, None] * (NUM_SPLITS * HEAD_DIM)
                + pid_s * HEAD_DIM
                + d[None, :])

    # online-softmax state (per padded q-head row)
    m_i = tl.full((PADDED_GROUP,), float("-inf"), dtype=tl.float32)
    l_i = tl.zeros((PADDED_GROUP,), dtype=tl.float32)
    acc = tl.zeros((PADDED_GROUP, HEAD_DIM), dtype=tl.float32)

    # load Q group: (PG, HD)
    q_ptrs = q + pid_b * (N_HEAD * HEAD_DIM) + qheads[:, None] * HEAD_DIM + d[None, :]
    q_tile = tl.load(q_ptrs, mask=row_mask[:, None], other=0.0)

    # this split's block range; empty splits skip the loop and write the
    # sentinel (m=-inf, l=0, acc=0) so the combine step ignores them.
    total_blocks = tl.cdiv(ctx, BLOCK_SIZE)
    blocks_per_split = tl.cdiv(total_blocks, NUM_SPLITS)
    start_block = pid_s * blocks_per_split
    end_block = tl.minimum(start_block + blocks_per_split, total_blocks)

    offs = tl.arange(0, BLOCK_SIZE)
    for blk in range(start_block, end_block):
        phys = tl.load(block_table + pid_b * MAX_BLOCKS_PER_SEQ + blk)
        phys = tl.maximum(phys, 0).to(tl.int64)
        base = phys * (N_KV_HEAD * BLOCK_SIZE * HEAD_DIM) + pid_kv * (BLOCK_SIZE * HEAD_DIM)

        tok_pos = blk * BLOCK_SIZE + offs                 # (BS,)
        kv_valid = tok_pos < ctx
        kv_ptrs = base + offs[:, None] * HEAD_DIM + d[None, :]   # (BS, HD)
        k_blk = tl.load(k_cache + kv_ptrs, mask=kv_valid[:, None], other=0.0)
        v_blk = tl.load(v_cache + kv_ptrs, mask=kv_valid[:, None], other=0.0)

        # S = Q @ K^T : (PG, BS)
        s = tl.dot(q_tile, tl.trans(k_blk)).to(tl.float32) * scale
        s = tl.where(kv_valid[None, :], s, float("-inf"))

        m_new = tl.maximum(m_i, tl.max(s, axis=1))
        alpha = tl.exp(m_i - m_new)
        p = tl.exp(s - m_new[:, None])                    # (PG, BS)
        l_i = l_i * alpha + tl.sum(p, axis=1)
        acc = acc * alpha[:, None] + tl.dot(p.to(v_blk.dtype), v_blk).to(tl.float32)
        m_i = m_new

    tl.store(pm_ptr, m_i, mask=row_mask)
    tl.store(pl_ptr, l_i, mask=row_mask)
    tl.store(pacc_ptr, acc, mask=row_mask[:, None])


@triton.jit
def _flash_decode_combine_kernel(
    partial_acc,        # (B, N_HEAD, NUM_SPLITS, HEAD_DIM) fp32
    partial_m,          # (B, N_HEAD, NUM_SPLITS) fp32
    partial_l,          # (B, N_HEAD, NUM_SPLITS) fp32
    context_lens,       # (B,) int32
    out,                # (B, N_HEAD, HEAD_DIM)
    HEAD_DIM: tl.constexpr,
    N_HEAD: tl.constexpr,
    NUM_SPLITS: tl.constexpr,
):
    pid_b = tl.program_id(0)
    pid_h = tl.program_id(1)
    d = tl.arange(0, HEAD_DIM)
    out_ptr = out + pid_b * (N_HEAD * HEAD_DIM) + pid_h * HEAD_DIM + d

    ctx = tl.load(context_lens + pid_b)
    if ctx == 0:
        tl.store(out_ptr, tl.zeros((HEAD_DIM,), dtype=out.dtype.element_ty))
        return

    s_idx = tl.arange(0, NUM_SPLITS)
    m_s = tl.load(partial_m + pid_b * (N_HEAD * NUM_SPLITS) + pid_h * NUM_SPLITS + s_idx)
    l_s = tl.load(partial_l + pid_b * (N_HEAD * NUM_SPLITS) + pid_h * NUM_SPLITS + s_idx)
    acc_ptrs = (partial_acc
                + pid_b * (N_HEAD * NUM_SPLITS * HEAD_DIM)
                + pid_h * (NUM_SPLITS * HEAD_DIM)
                + s_idx[:, None] * HEAD_DIM
                + d[None, :])
    acc_s = tl.load(acc_ptrs)                            # (NS, HD)

    m_g = tl.max(m_s)                                    # scalar
    factor = tl.exp(m_s - m_g)                           # (NS,)
    l_g = tl.sum(l_s * factor)
    acc_g = tl.sum(acc_s * factor[:, None], axis=0)      # (HD,)
    out_v = acc_g / l_g
    tl.store(out_ptr, out_v.to(out.dtype.element_ty))