CS336-Assignment2

# 用法举例
@nvtx.range("scaled dot product attention") # 装饰器一个函数测量

with nvtx.range("computing attention scores"): # 测量一段代码时间
    attention = q @ k.transpose(-1,-2) / d_k ** 0.5

# push 和pop中间是所测量的部分
nvtx.range_push(f"forward pass test {i-4}")
forward_pass_time.append(timeit.timeit(forward_pass_only, number=1))
nvtx.range_pop()

            # O_i = tl.exp(m_i_old - m_i)[:, None] * O_i + tl.dot(P_i, V_j)
            # 优化版本（使用 acc）
            O_i = tl.exp(m_i_old - m_i)[:, None] * O_i
            O_i = tl.dot(P_i, V_j, acc=O_i)  # 累积到 O_i

    # 将 O_i 转换回原始数据类型（与 Q 相同）,使用*_block_ptr.type.element_ty
    tl.store(O_block_ptr, O_i.to(O_block_ptr.type.element_ty), boundary_check=(0, 1))

# Kernel 3: Forward Pass
@triton.jit
def flash_fwd_kernel(
    Q_ptr, K_ptr, V_ptr, # 输入指针
    O_ptr, L_ptr,# 输出指针，L式logsumexp需要保存用于反向传播
    stride_qb, stride_qq, stride_qd, # Q的stride，包括跳一个batch，跳一个query，跳一个dimension所需要的步长
    stride_kb, stride_kk, stride_kd,
    stride_vb, stride_vk, stride_vd,
    stride_ob, stride_oq, stride_od,
    stride_lb, stride_lq, # L的stride，二维tensor
    N_QUERIES, N_KEYS, # 一般等于seq_len
    scale, # 1/sqrt(d)
    D: tl.constexpr, # d维度大小
    Q_TILE_SIZE: tl.constexpr, # 每次处理的query块大小
    K_TILE_SIZE: tl.constexpr, # 每次处理的key value块大小
    is_causal: tl.constexpr # 是否是causal attention，需要掩码
):
    query_tile_index = tl.program_id(0) # 第几个query块，第0维度并行性更好，用来划分query块
    batch_index = tl.program_id(1) # 第几个batch，第1维度并行性较差，用来划分batch
    #     tl.make_block_ptr(
    #     base,         # 基地址（指针）
    #     shape,        # 整个张量的形状
    #     strides,      # 步长（stride）
    #     offsets,      # 块的起始位置
    #     block_shape,  # 块的大小
    #     order         # 内存布局顺序
    # )
    Q_block_ptr = tl.make_block_ptr(
        Q_ptr + batch_index * stride_qb, # batch维度偏移，这个batch的起始位置
        shape = (N_QUERIES, D), # Q的整体shape，即(seq_len, d),也是（N_QUERIES, D)，用于边界检查
        strides = (stride_qq, stride_qd),
        offsets = (query_tile_index*Q_TILE_SIZE, 0), # 每个query块的起始位置
        block_shape = (Q_TILE_SIZE, D), # 每个块的shape
        order = (1, 0), # 内存布局顺序，遍历顺序，这里是d维度优先
    )
    K_block_ptr = tl.make_block_ptr(
        K_ptr + batch_index * stride_kb,
        shape = (N_KEYS, D),
        strides = (stride_kk, stride_kd),
        offsets = (0, 0), # K, V的offsets都应该从0开始，随j遍历是内循环
        block_shape = (K_TILE_SIZE, D),
        order = (1, 0),
    )
    V_block_ptr = tl.make_block_ptr(
        V_ptr + batch_index * stride_vb,
        shape = (N_KEYS, D),
        strides = (stride_vk, stride_vd),
        offsets = (0, 0),
        block_shape = (K_TILE_SIZE, D),
        order = (1, 0),
    )
    O_block_ptr = tl.make_block_ptr(
        O_ptr + batch_index * stride_ob,
        shape = (N_QUERIES, D),
        strides = (stride_oq, stride_od),
        offsets = (query_tile_index*Q_TILE_SIZE, 0),
        block_shape = (Q_TILE_SIZE, D),
        order = (1, 0),
    )
    L_block_ptr = tl.make_block_ptr(
        L_ptr + batch_index * stride_lb,
        shape = (N_QUERIES, ), # L是一维tensor，shape是(seq_len,)
        strides = (stride_lq, ),
        offsets = (query_tile_index*Q_TILE_SIZE,),
        block_shape = (Q_TILE_SIZE,),
        order = (0,), # 一维tensor只有一个维度
    )
    O_i = tl.zeros((Q_TILE_SIZE, D), dtype=tl.float32) # 需要初始化，注o数据类型
    l_i = tl.zeros((Q_TILE_SIZE,), dtype=tl.float32)
    m_i = tl.full((Q_TILE_SIZE,),  -float("inf"), dtype=tl.float32) # tl.full的使用，-float("inf")

    #  加载Q块，有边界检查
    Q_i = tl.load(Q_block_ptr, boundary_check=(0, 1), padding_option="zero")
    Q_i = tl.cast(Q_i, tl.float32) # 数据类型转换

    for j in range(tl.cdiv(N_KEYS, K_TILE_SIZE)):
        # j_start表示当前块在key维度的起始位置
        j_start = j * K_TILE_SIZE  
        # 判断是否需要跳过这个块
        should_skip = False
        if is_causal:
            q_max_idx = (query_tile_index + 1) * Q_TILE_SIZE - 1
            k_min_idx = j_start
            should_skip = q_max_idx < k_min_idx
        if not should_skip:
            # 创建 key index 向量并判断哪些是有效（< N_KEYS）
            # 这两行代码的作用是创建有效性掩码（validity mask），用于处理序列长度不能被块大小整除的边界情况
            # 具体来说：
            # j=3, j_start = 3 * 32 = 96
            # k_idx = tl.arange(0, 32) + 96
            # k_idx = [96, 97, 98, 99, 100, 101, ..., 127]
            # shape: (32,)

            # N_KEYS = 100
            # valid_k = k_idx < 100
            # valid_k = [True, True, True, True, False, False, ..., False]
            #           (96-99 为 True，100-127 为 False)
            # shape: (32,)
            k_idx = tl.arange(0, K_TILE_SIZE) + j_start  # shape (K_TILE_SIZE,)
            valid_k = k_idx < N_KEYS                      # boolean mask shape (K_TILE_SIZE,)
            K_j = tl.load(K_block_ptr, boundary_check=(0, 1), padding_option="zero")
            V_j = tl.load(V_block_ptr, boundary_check=(0, 1), padding_option="zero")
            K_j = tl.cast(K_j, tl.float32) # 数据类型
            V_j = tl.cast(V_j, tl.float32)
            S_i = tl.dot(Q_i, tl.trans(K_j)) * scale
            # mask operation 
            if is_causal:
                # 计算块中元素在S中的的行数
                q_idx = tl.arange(0, Q_TILE_SIZE) + query_tile_index * Q_TILE_SIZE  # shape (Q_TILE_SIZE,)
                # 计算块中元素在S中的的列数 
                k_idx_tile = tl.arange(0, K_TILE_SIZE) + j_start  # shape (K_TILE_SIZE,)
                # 行数大于列数就被mask，合理利用广播机制
                causal_mask = q_idx[:, None] >= k_idx_tile[None, :]  # shape (Q_TILE_SIZE, K_TILE_SIZE)
                #        Key 位置 →
                #         k4  k5  k6  k7
                # Query q4 [T,  F,  F,  F]  ← query 4 只能看到 key 4 (自己)
                # 位置  q5 [T,  T,  F,  F]  ← query 5 可以看到 key 4-5
                # ↓     q6 [T,  T,  T,  F]  ← query 6 可以看到 key 4-6
                #       q7 [T,  T,  T,  T]  ← query 7 可以看到 key 4-7

                # 解释：
                # q4 >= k4 → T,  q4 >= k5 → F,  q4 >= k6 → F,  q4 >= k7 → F
                # q5 >= k4 → T,  q5 >= k5 → T,  q5 >= k6 → F,  q5 >= k7 → F
                # q6 >= k4 → T,  q6 >= k5 → T,  q6 >= k6 → T,  q6 >= k7 → F
                # q7 >= k4 → T,  q7 >= k5 → T,  q7 >= k6 → T,  q7 >= k7 → T
                # Apply causal mask: add -1e6 to masked out elements
                S_i = tl.where(causal_mask, S_i, S_i - 1e6)
                # tl.where(condition, if_true, if_false)
                # 如果 causal_mask[i, j] == True:  保持 S_i[i, j]
                # 如果 causal_mask[i, j] == False: S_i[i, j] = S_i[i, j] - 1e6
            
            # 广播应用：整列掩码
            # valid_k[None, :] -> [1, K_TILE_SIZE]
            #        k96 k97 k98 k99 k100 k101 ... k127
            # q0  [[ T,  T,  T,  T,  F,   F,  ... F],   # 每列统一
            # q1   [ T,  T,  T,  T,  F,   F,  ... F],
            # q2   [ T,  T,  T,  T,  F,   F,  ... F],
            # q3   [ T,  T,  T,  T,  F,   F,  ... F]]
            S_i = tl.where(valid_k[None,:], S_i, -float("inf"))
            m_i_old = m_i
            m_i = tl.maximum(m_i_old, tl.max(S_i, axis=-1))
            P_i = tl.exp(S_i - m_i[:, None]) # triton中不能用...,None
            l_i = tl.exp(m_i_old - m_i) * l_i + tl.sum(P_i, axis=-1)
            # O_i = tl.exp(m_i_old - m_i)[:, None] * O_i + tl.dot(P_i, V_j)
            # 优化版本（使用 acc）
            O_i = tl.exp(m_i_old - m_i)[:, None] * O_i
            O_i = tl.dot(P_i, V_j, acc=O_i)  # 累积到 O_i
        K_block_ptr = K_block_ptr.advance((K_TILE_SIZE, 0)) # 相当于移动offset
        V_block_ptr = V_block_ptr.advance((K_TILE_SIZE, 0))
    O_i = O_i / l_i[:, None]
    L_i = m_i + tl.log(l_i) # 真实的logsumexp值
    # 将 O_i 转换回原始数据类型（与 Q 相同）
    tl.store(O_block_ptr, O_i.to(O_block_ptr.type.element_ty), boundary_check=(0, 1))
    tl.store(L_block_ptr, L_i.to(L_block_ptr.type.element_ty), boundary_check=(0,))

    for j in range(tl.cdiv(N_KEYS, K_TILE_SIZE)):
        j_start = j * K_TILE_SIZE
        should_skip = False
        if is_causal:
            q_max_idx = (query_tile_index + 1) * Q_TILE_SIZE - 1
            k_min_idx = j_start
            should_skip = q_max_idx < k_min_idx
        
        if not should_skip:  

# 定义 autotune 的配置搜索空间
# 这会尝试不同的块大小、warp数量和流水线级数，找到性能最好的组合
def get_configs():
    return [
        # --- 方案A: 针对 bfloat16 + d<=64 的高性能配置 ---
        # 显存消耗: 高 (stages=3 需要大量 SRAM)
        triton.Config({'Q_TILE_SIZE': 128, 'K_TILE_SIZE': 64}, num_warps=8, num_stages=3),
        triton.Config({'Q_TILE_SIZE': 64,  'K_TILE_SIZE': 64}, num_warps=4, num_stages=3),
        triton.Config({'Q_TILE_SIZE': 128, 'K_TILE_SIZE': 32}, num_warps=4, num_stages=3),

        # --- 方案B: 针对 bfloat16 + d=128 或 float32 + small d 的平衡配置 ---
        # 显存消耗: 中 (stages=2 节省显存)
        triton.Config({'Q_TILE_SIZE': 128, 'K_TILE_SIZE': 64}, num_warps=8, num_stages=2),
        triton.Config({'Q_TILE_SIZE': 64,  'K_TILE_SIZE': 64}, num_warps=4, num_stages=2),
        triton.Config({'Q_TILE_SIZE': 128, 'K_TILE_SIZE': 32}, num_warps=4, num_stages=2),
        triton.Config({'Q_TILE_SIZE': 64,  'K_TILE_SIZE': 32}, num_warps=4, num_stages=2),
        triton.Config({'Q_TILE_SIZE': 32,  'K_TILE_SIZE': 32}, num_warps=4, num_stages=2),
        
        # --- 方案C: 针对 float32 + d=128 的"求生"配置 ---
        # 显存消耗: 低 (stages=1 放弃流水线，只求不爆显存)
        # 这里的 32x32 + stage=1 几乎可以适应任何苛刻环境
        triton.Config({'Q_TILE_SIZE': 32, 'K_TILE_SIZE': 32}, num_warps=4, num_stages=1),
        triton.Config({'Q_TILE_SIZE': 64, 'K_TILE_SIZE': 32}, num_warps=4, num_stages=1),
    ]

# -----------------------------------------------------------------------------
# Kernel 1: 计算 dQ
@triton.autotune(
    configs=get_configs(),
    key=['N_QUERIES', 'N_KEYS'], # 当 seq_len 变化时重新调优
)
····
        # [修改点]：Grid 使用 lambda，根据 Autotune 选出的 META['Q_TILE_SIZE'] 动态计算
        grid = lambda META: (triton.cdiv(N_QUERIES, META['Q_TILE_SIZE']), batch_size)
        flash_fwd_kernel[grid](

import torch.distributed as dist

def setup(rank, world_size, backend):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29100"
    dist.init_process_group(backend=backend,rank=rank,world_size=world_size)
    
dist.all_reduce(parameter.grad, op=dist.ReduceOp.SUM, async_op=True)
# 确保所有进程同步后再清理
dist.barrier()
# 清理分布式进程组
dist.destroy_process_group()
# 获取rank world-size
self.rank = dist.get_rank()
self.world_size = dist.get_world_size()

        flattened_grads = torch._utils._flatten_dense_tensors(grads)
        dist.all_reduce(flattened_grads, op=dist.ReduceOp.AVG,async_op=False)
        unflattened_grad = torch._utils._unflatten_dense_tensors(flattened_grads, grads)

# 代码片段：ddp_overlap_bucketed.py
def _grad_hook(self, param):
    bucket_idx = self.param_to_bucket[param]
    # 检查桶里的梯度是否都就绪了
    if self._is_bucket_ready(bucket):
        # 拍扁 -> 异步通信
        self._start_bucket_communication(bucket_idx)

class ShardedOptimizer(torch.optim.Optimizer):
    def __init__(self, params, optimizer, **kwargs):
        self.rank = dist.get_rank()
        self.world_size = dist.get_world_size()
        super().__init__(params, kwargs)
        self.optimizer = optimizer(self.local_param_group['params'], **kwargs) # 先准备好local_param_group，再实例化本地的optimizer
        
    @torch.no_grad()
    def step(self,closure=None, **kwargs):# 注意参数传递，这个closure
        self.optimizer.step(closure, **kwargs)
        for src_rank in range(self.world_size): # 一定注意broadcase的方法，源是谁，接收是谁
            for param_group in self.param_groups:
                for i, param in enumerate(param_group['params']):
                    if i % self.world_size == src_rank:
                        dist.broadcast(param.data,src=src_rank)
    
    def add_param_group(self, param_group: Dict[str, Any]):
        # 传入的 param_group 是一个字典，其中 'params' 键对应的是一个参数列表。我们需要遍历这个列表。
        # self.param_groups: List[Dict[str, Any]] = []
        # 需要考虑到传入的是多个param_group的情况
        super().add_param_group(param_group) # 先保存一份完整的模型参数
        local_param_group: Dict[str, Any] = {'params':[]} # 只需要存param，不需要存其他的，因为实例化的时候有**kwargs
        for i, param in enumerate(param_group['params']):
            if i % self.world_size == self.rank:
                local_param_group['params'].append(param)
        self.local_param_group = local_param_group