import torch

N, d = 1024, 64  # Обновить значение Nид

Q_mat = torch.rand((N, d))
K_mat = torch.rand((N, d))
V_mat = torch.rand((N, d))

def standard_softmax_attention(Q, K, V):
    """
    Внедрить стандартный Pytorch softmax расчет внимания.
    """
    expected_softmax = torch.softmax(Q @ K.T, dim=1)
    expected_attention = expected_softmax @ V
    return expected_softmax, expected_attention

def safe_softmax_attention(Q, K, V):
    """
    执ХОРОШО安全изsoftmax расчет внимания.
    """
    S_mat = Q @ K.T
    row_max = torch.max(S_mat, dim=1).values[:, None]
    input_safe = S_mat - row_max
    softmax_numerator = torch.exp(input_safe)
    softmax_denominator = torch.sum(softmax_numerator, dim=1)[:, None]
    safe_softmax = softmax_numerator / softmax_denominator
    matmul_result = safe_softmax @ V
    return safe_softmax, matmul_result

# Рассчитано с использованием стандартного softmaxиattention.
expected_softmax, expected_attention = standard_softmax_attention(Q_mat, K_mat, V_mat)
# Вычисление с использованием безопасного softmaxвнимания
safe_softmax, safe_attention = safe_softmax_attention(Q_mat, K_mat, V_mat)

# Убедитесь, что результаты softmax и внимания, рассчитанные двумя методами, близки.
assert torch.allclose(safe_softmax, expected_softmax), "error in safe softmax"
assert torch.allclose(safe_attention, expected_attention), "error in safe attention"

import torch

N, d = 1024, 64  # Обновить значение Nид

Q_mat = torch.rand((N, d))
K_mat = torch.rand((N, d))
V_mat = torch.rand((N, d))

def standard_softmax_attention(Q, K, V):
    """
    Внедрить стандартный Pytorch softmax расчет внимания.
    """
    expected_softmax = torch.softmax(Q @ K.T, dim=1)
    expected_attention = expected_softmax @ V
    return expected_softmax, expected_attention

def flash_attention(Q, K, V, B_r=64, B_c=768):
    """
    Выполнить прошивку с использованием блочного расчета исуществовать коррекцию softmax строки алгоритм внимания.
    """
    O = torch.zeros((N, d))  # Инициализируйте выходную матрицу, соответствующую строке 2 псевдокода.
    l = torch.zeros((N, 1))  # Сохраните знаменатель softmax, соответствующий строке 2 псевдокода.
    m = torch.full((N, 1), -torch.inf)  # Сохраните максимальное значение каждого блока, соответствующее второй строке псевдокода.

    # Соответствует строке 5 псевдокода, для 1<=j<=T_c,Обратите внимание, что здесьK, V разделен на блоки T_c=[N/B_c], размер каждого блока равен [B_c, г] Такой большой
    # Поэтому, когда реализуется существованиеpython, он обрабатывается напрямую в цикле с размером шага B_c.
    for j in range(0, N, B_c):
        # Следующие три строки соответствуют строке 6 псевдокода Load Kj, Vj from HBM to on-chip SRAM
        # Но вот простота python Реализация, у нас невозможно реально поставить этот кусок памяти из HBM в SRAM
        # Это всего лишь логическое объяснение псевдокода, вы можете сделать вид, что оно это делает, потому что существующий Triton действительно может делать это в существующем слое Python.
        j_end = j + B_c
        Kj = K[j:j_end, :]
        Vj = V[j:j_end, :]

        # Соответствует строке 7 псевдокода, для 1<=i<T_r,Обратите внимание, что Q разделен на блоки Tr=[N/B_r],每一块из大小да[B_r, г] Такой большой
        # Поэтому, когда реализуется существованиеpython, он обрабатывается напрямую в цикле с размером шага B_r.
        for i in range(0, N, B_r):
            i_end = i + B_r
            mi = m[i:i_end, :]
            li = l[i:i_end, :]
            Oi = O[i:i_end, :]
            Qi = Q[i:i_end, :]

            # Соответствует строке 9 псевдокода: on chip, compute Sij, форма Sij — [B_r, B_c]
            Sij = Qi @ Kj.T
            # Соответствует строке 10 псевдокода
            mij_hat = torch.max(Sij, dim=1).values[:, None]
            pij_hat = torch.exp(Sij - mij_hat)
            lij_hat = torch.sum(pij_hat, dim=1)[:, None]

            # В соответствии со строкой 11 псевдокода находится операция mi_new. Обратите внимание, что здесь необходимо найти общий максимум двух тензоров, поэтому существует эта операция стека.
            mi_new = torch.max(torch.column_stack([mi, mij_hat]), dim=1).values[:, None]
            # По строке 11 псевдокода найдите операцию li_new
            li_new = torch.exp(mi - mi_new) * li + torch.exp(mij_hat - mi_new) * lij_hat
            # В соответствии со строкой 12 псевдокода O_i обновляется. Здесь возникает вопрос. В псевдокоде есть операция диагностики. Почему она игнорируется в следующей реализации?
            # Это связано с тем, что эта диаграмма используется для существующего вектора. На самом деле она соответствует верхнему измерению в существующем псевдокоде, а реализация PyTorch является автоматической.
            # Поддерживается механизм тензорной трансляции, поэтому его можно рассчитать прямо здесь.
            O_i = (li * torch.exp(mi - mi_new) * Oi / li_new) + (torch.exp(mij_hat - mi_new) * pij_hat / li_new) @ Vj

            # В соответствии со строкой 13 псевдокода обновляются m_i, l_i, O_i.
            m[i:i_end, :] = mi_new
            l[i:i_end, :] = li_new
            O[i:i_end, :] = O_i

    return O

# выполнить флэш расчет внимания
flash_attention_output = flash_attention(Q_mat, K_mat, V_mat)

# Внедрить стандартный Pytorch softmaxирасчет внимания
expected_softmax, expected_attention = standard_softmax_attention(Q_mat, K_mat, V_mat)

# вспышка подтверждения расчет вниманияиз结果与标准вычислить结果да否接近
assert torch.allclose(flash_attention_output, expected_attention), "error in flash attention calculation"

import torch

N, d = 1024, 64  # Обновить значение Nид

Q_mat = torch.rand((N, d))
K_mat = torch.rand((N, d))
V_mat = torch.rand((N, d))

def standard_softmax_attention(Q, K, V):
    """
    Реализация стандартного PyTorch softmax расчет внимания.
    """
    expected_softmax = torch.softmax(Q @ K.T, dim=1)
    expected_attention = expected_softmax @ V
    return expected_softmax, expected_attention

def flash_attention_v2(Q, K, V, B_r=64, B_c=768):
    """
    Выполнить прошивку с использованием блочного расчета исуществовать коррекцию softmax строки attention алгоритм v2.
    """
    O = torch.zeros((N, d))  # Инициализируйте O как (N, Форма d) фактически соответствует инициализации O в строке 5 псевдокода.
    l = torch.zeros((N, 1))  # Инициализация l в форме (N) фактически соответствует инициализации l в строке 5 псевдокода.
    m = torch.full((N, 1), -torch.inf)  # Сохраните максимальное значение каждого блока, инициализированное отрицательной бесконечностью, соответствующее строке 5 псевдокода.

    # Соответствует строке 3 псевдокода, для 1<=i<T_r,Обратите внимание, что Q разделен на блоки Tr=[N/B_r],每一块из大小да[B_r, г] Такой большой
    # Поэтому, когда реализуется существованиеpython, он обрабатывается напрямую в цикле с размером шага B_r.
    for i in range(0, N, B_r):
        Qi = Q[i:i+B_r, :]
        # Соответствует псевдокоду 6 ОК, для 1<=j<=T_c,Обратите внимание, что здесьK, V разделен на блоки T_c=[N/B_c], размер каждого блока равен [B_c, г] Такой большой
        # Поэтому, когда реализуется существованиеpython, он обрабатывается напрямую в цикле с размером шага B_c. 
        for j in range(0, N, B_c):  # Внутренний цикл пересекает блоки Q
            Kj = K[j:j+B_c, :]
            Vj = V[j:j+B_c, :]

            # Соответствует псевдокоду8ХОРОШО：on chip, compute Sij, форма Sij — [B_r, B_c]
            Sij = Qi @ Kj.T
            # Соответствует псевдокоду9ХОРОШОпроситьm_i^(j)из操作,mi_newизформадаB_r
            mi_new = torch.max(torch.column_stack([m[i:i+B_r], torch.max(Sij, dim=1).values[:, None]]), dim=1).values[:, None]
            # Соответствует псевдокоду9ХОРОШОпроситьPij_hatиз操作,Форма Пидж_хата равна (B_r x B_c)., и Sij непротиворечив
            Pij_hat = torch.exp(Sij - mi_new)
            # Соответствует псевдокоду9ХОРОШОпроситьlijиз操作
            l[i:i+B_r] = torch.exp(m[i:i+B_r] - mi_new) * l[i:i+B_r] + torch.sum(Pij_hat, dim=1)[:, None]
            # Соответствует строке 10 псевдокодапроситьO_ijиз操作
            O[i:i+B_r] = O[i:i+B_r] * torch.exp(m[i:i+B_r] - mi_new) + Pij_hat @ Vj
            m[i:i+B_r] = mi_new

    O = O / l  # В соответствии со строкой 12 псевдокода выходные данные корректируются в соответствии со знаменателем softmax.
    return O

# выполнить флэш расчет внимания
flash_attention_v2_output = flash_attention_v2(Q_mat, K_mat, V_mat)

# Реализация стандартного PyTorch softmaxирасчет внимания
_, expected_attention = standard_softmax_attention(Q_mat, K_mat, V_mat)

# вспышка подтверждения расчет вниманияиз结果与标准вычислить结果да否接近
assert torch.allclose(flash_attention_v2_output, expected_attention), "Error in flash attention calculation"

@triton.jit
def _attn_fwd_inner(acc, l_i, m_i, q,  #
                    K_block_ptr, V_block_ptr,  #
                    start_m, qk_scale,  #
                    BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,  #
                    STAGE: tl.constexpr, offs_m: tl.constexpr, offs_n: tl.constexpr,  #
                    N_CTX: tl.constexpr):
    # range of values handled by this stage
    # На основе значения STAGE функция определяет диапазон ключей (K) и значений (V), подлежащих обработке.
    # Разные ЭТАПЫ соответствуют разным диапазонам обработки, поддерживая причинное и непричинное внимание к себе.
    if STAGE == 1: # causal = True
        lo, hi = 0, start_m * BLOCK_M
    elif STAGE == 2:
        lo, hi = start_m * BLOCK_M, (start_m + 1) * BLOCK_M
        lo = tl.multiple_of(lo, BLOCK_M)
    # causal = False
    else:
        lo, hi = 0, N_CTX
    # Используйте функцию tl.advance для настройки положения указателя KиV для корректной загрузки данных из соответствующей ячейки памяти.
    K_block_ptr = tl.advance(K_block_ptr, (0, lo))
    V_block_ptr = tl.advance(V_block_ptr, (lo, 0))
    # loop over k, v and update accumulator
    # существуют В цикле функция загружает блок ключей (K), вычисляет скалярное произведение запроса (Q) и этого блока ключей,
    # Затем скорректируйте результаты расчета на основе текущего ЭТАПА. Если это ЭТАП 2 и причинно-следственная связь верна, маска применяется для блокировки будущей информации.
    for start_n in range(lo, hi, BLOCK_N):
        start_n = tl.multiple_of(start_n, BLOCK_N)
        # -- compute qk ----
        k = tl.load(K_block_ptr)
        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
        qk += tl.dot(q, k)
        if STAGE == 2:
            mask = offs_m[:, None] >= (start_n + offs_n[None, :])
            qk = qk * qk_scale + tl.where(mask, 0, -1.0e6)
            m_ij = tl.maximum(m_i, tl.max(qk, 1))
            qk -= m_ij[:, None]
        else:
            # В соответствии с вычислением m_ij в строке 9 псевдокода потока алгоритма, разница между псевдокодом и псевдокодом заключается в том, что здесь применяется qk_scale.
            m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
            qk = qk * qk_scale - m_ij[:, None]
        # Рассчитать р,Соответствует псевдокоду9ХОРОШОизpизвычислить
        p = tl.math.exp2(qk)
        l_ij = tl.sum(p, 1)
        # -- update m_i and l_i
        alpha = tl.math.exp2(m_i - m_ij)
        l_i = l_i * alpha + l_ij
        # -- update output accumulator --
        acc = acc * alpha[:, None]
        # update acc
        v = tl.load(V_block_ptr)
        acc += tl.dot(p.to(tl.float16), v)
        # update m_i and l_i
        m_i = m_ij
        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
    return acc, l_i, m_i

# Определен класс _attention, который наследуется от torch.autograd.Function. Это позволяет нам настраивать прямое и обратное распространение операции.
# (то есть способ расчета градиентов), что позволяет легко интегрировать его с системой автоматического расчета градиентов PyTorch.
class _attention(torch.autograd.Function):

    @staticmethod
    # Метод Forward определяет логику прямого распространения этой пользовательской операции. ctx — это объект контекста, используемый для хранения информации, используемой для обратного распространения ошибки.
    # q, k, v соответственно представляет запрос, key, value имеет три входных тензора, а causalism_scale — дополнительный управляющий параметр.
    def forward(ctx, q, k, v, causal, sm_scale):
        # shape constraints
        # Эти строки кода проверяют последнее измерение входного тензора, чтобы убедиться, что они имеют одинаковый размер и определенное значение (16, 32, 64, или 128). Это связано с необходимостью оптимизации производительности для конкретной реализации.
        Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
        assert Lq == Lk and Lk == Lv
        assert Lk in {16, 32, 64, 128}
        # Инициализируйте пустой Tensoro той же формы и типа, что и q, чтобы сохранить выходные результаты.
        o = torch.empty_like(q)
        # Эти строки устанавливают несколько ключевых параметров настройки производительности, включая размер блока обработки (BLOCK_M, BLOCK_N）и
        # Подсчитайте количество этапов (num_stages). num_warps относится к каждому CUDA Количество варпов в блоке.
        BLOCK_M = 128
        BLOCK_N = 64 if Lk <= 64 else 32
        num_stages = 4 if Lk <= 64 else 3
        num_warps = 4
        stage = 3 if causal else 1
        # Согласно возможностям устройства CUDA (здесь проверяются вычислительные возможности 9.x, т.е. NVIDIA Архитектура Volta и более поздние архитектуры), далее настройте num_warpsиnum_stages.
        # Tuning for H100
        if torch.cuda.get_device_capability()[0] == 9:
            num_warps = 8
            num_stages = 7 if Lk >= 64 else 3
        # Рассчитать Тритон Размер сетки ядра. triton.cdiv — это вспомогательная функция, которая вычисляет деление путем округления в большую сторону.
        # q.shape[2] — длина последовательности, q.shape[0] и q.shape[1] — партия и последовательность соответственно. length
        grid = (triton.cdiv(q.shape[2], BLOCK_M), q.shape[0] * q.shape[1], 1)
        # Инициализируйте еще один TensorM, который используется для хранения промежуточных результатов в процессе расчета.
        M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
        # Позвоните Тритону kernel _attn_fwd执ХОРОШО实际израсчет внимания。这里传递了大量изпараметр,Включить вводTensorиз各个维度、шагать、форма、Параметры настройки и т. д..
        _attn_fwd[сетка](
            q, k, v, sm_scale, M, o, #
            q.шаг(0), q.шаг(1), q.шаг(2), q.шаг(3), #
            k.stride(0), k.stride(1), k.stride(2), k.stride(3), #
            v.stride(0), v.stride(1), v.stride(2), v.stride(3), #
            o.stride(0), o.stride(1), o.stride(2), o.stride(3), #
            q.shape[0], q.shape[1], #
            N_CTX=q.shape[2], #
            БЛОК_М=БЛОК_М, #
            БЛОК_Н=БЛОК_Н, #
            BLOCK_DMODEL=Лк, #
            ЭТАП = этап, #
            num_warps=num_warps, #
            num_stages=num_stages #
        )

BLOCK_M = 128
BLOCK_N = 64 if Lk <= 64 else 32

@triton.jit
# Определена функция _attn_fwd. Эта функция представляет собой ядро, реализующее прямой проход механизма внимания. Параметры функции включают входные тензоры запроса (Q), ключа (K) и значения (V).
# коэффициент масштабирования softmax (sm_scale), промежуточный результат расчета (M) и выходной тензор (Out), а Также несколько параметров шага и другие константы конфигурации для этих тензоров.
def _attn_fwd(Q, K, V, sm_scale, M, Out,  #
              stride_qz, stride_qh, stride_qm, stride_qk,  #
              stride_kz, stride_kh, stride_kn, stride_kk,  #
              stride_vz, stride_vh, stride_vk, stride_vn,  #
              stride_oz, stride_oh, stride_om, stride_on,  #
              Z, H,  #
              N_CTX: tl.constexpr,  #
              BLOCK_M: tl.constexpr,  #
              BLOCK_DMODEL: tl.constexpr,  #
              BLOCK_N: tl.constexpr,  #
              STAGE: tl.constexpr  #
              ):
    # Обратите внимание, что ZиH во входных параметрах означает партию соответственно. размерколичество голов внимания
    # start_m представляет текущее ядро program Смещение измерения seq, соответствующего экземпляру, а off_hz представляет собой смещение измерения партия*головки.
    start_m = tl.program_id(0)
    off_hz = tl.program_id(1)

    # 这些ХОРОШОвычислить了两个偏移量off_zиoff_h,Они соответственно представляютсуществоватьbatch（илиheads）中из位置。
    off_z = off_hz // H
    off_h = off_hz % H
    # Вычислите смещение, используемое для поиска текущего блока обработки в тензорах Q, KиV. Это основано на ранее вычисленном смещении и предоставленном параметре шага.
    qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh

    # block pointers
    # Используйте tl.make_block_ptr, чтобы создать указатель на текущий блок обработки Q-тензора. Этот вызов функции определяет базовый адрес, форму, шаг, смещение, форму блока и т. д. также Как получить доступ к этому блоку данных в памяти.
    # N_CTX — это q.shape[2], который представляет длину последовательности, а BLOCK_DMODEL — это Lk, который представляет размерность скрытого слоя каждой головки внимания.
    # Следующие тензоры, созданные make_block_ptr, аналогичны, а именно для K, Vа ТакжеOutputO создает указатель на текущий блок обработки.
    Q_block_ptr = tl.make_block_ptr(
        base=Q + qvk_offset,
        shape=(N_CTX, BLOCK_DMODEL),
        strides=(stride_qm, stride_qk),
        offsets=(start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, BLOCK_DMODEL),
        order=(1, 0),
    )
    V_block_ptr = tl.make_block_ptr(
        base=V + qvk_offset,
        shape=(N_CTX, BLOCK_DMODEL),
        strides=(stride_vk, stride_vn),
        offsets=(0, 0),
        block_shape=(BLOCK_N, BLOCK_DMODEL),
        order=(1, 0),
    )
    K_block_ptr = tl.make_block_ptr(
        base=K + qvk_offset,
        shape=(BLOCK_DMODEL, N_CTX),
        strides=(stride_kk, stride_kn),
        offsets=(0, 0),
        block_shape=(BLOCK_DMODEL, BLOCK_N),
        order=(0, 1),
    )
    O_block_ptr = tl.make_block_ptr(
        base=Out + qvk_offset,
        shape=(N_CTX, BLOCK_DMODEL),
        strides=(stride_om, stride_on),
        offsets=(start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, BLOCK_DMODEL),
        order=(1, 0),
    )
    # initialize offsets
    # Вычислите начальное смещение элементов, которые должен обрабатывать каждый поток, в измерении M (последовательное измерение).
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    # Вычислите смещение элементов, которые должен обрабатывать каждый поток, в измерении N (размерность партии*головок).
    offs_n = tl.arange(0, BLOCK_N)
    # initialize pointer to m and l
    # Инициализируйте вектор m, m используется для хранения максимального логита в каждом измерении m и инициализируется отрицательной бесконечностью.
    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
    # Инициализируйте вектор l, l используется как знаменатель совокупного softmax и инициализируется значением 1.
    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
    # Инициализируйте аккумулятор для накопления весов внимания. Обратите внимание, что здесь фигура (BLOCK_M, BLOCK_DMODEL)
    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
    # load scales
    qk_scale = sm_scale     # Загрузите коэффициент масштабирования softmax.
    qk_scale *= 1.44269504  # Умножьте коэффициент масштабирования softmax на 1/log(2) для последующих вычислений.
    # load q: it will stay in SRAM throughout
    q = tl.load(Q_block_ptr) # Загрузите текущий блок Q-матрицы в SRAM, эти данные остаются неизменными на протяжении всего расчета.
    # stage 1: off-band
    # For causal = True, STAGE = 3 and _attn_fwd_inner gets 1 as its STAGE
    # For causal = False, STAGE = 1, and _attn_fwd_inner gets 3 as its STAGE
    if STAGE & 1:
        acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, K_block_ptr, V_block_ptr,  #
                                        start_m, qk_scale,  #
                                        BLOCK_M, BLOCK_DMODEL, BLOCK_N,  #
                                        4 - STAGE, offs_m, offs_n, N_CTX  #
                                        )
    # stage 2: on-band
    if STAGE & 2:
        # barrier makes it easier for compielr to schedule the
        # two loops independently
        tl.debug_barrier()
        acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, K_block_ptr, V_block_ptr,  #
                                        start_m, qk_scale,  #
                                        BLOCK_M, BLOCK_DMODEL, BLOCK_N,  #
                                        2, offs_m, offs_n, N_CTX  #
                                        )
    # epilogue
    m_i += tl.math.log2(l_i)
    acc = acc / l_i[:, None]
    m_ptrs = M + off_hz * N_CTX + offs_m
    tl.store(m_ptrs, m_i)
    tl.store(O_block_ptr, acc.to(Out.type.element_ty))