Initial commit

flash_attn_triton.py

@@ -0,0 +1,861 @@
+"""
+Copied from https://github.com/HazyResearch/flash-attention/blob/eff9fe6b8076df59d64d7a3f464696738a3c7c24/flash_attn/flash_attn_triton.py
+update imports to use 'triton_pre_mlir'
+
+*Experimental* implementation of FlashAttention in Triton.
+Tested with triton==2.0.0.dev20221202.
+Triton 2.0 has a new backend (MLIR) but seems like it doesn't yet work for head dimensions
+other than 64:
+https://github.com/openai/triton/blob/d376020f90002757eea3ea9475d4f7cfc2ec5ead/python/triton/ops/flash_attention.py#L207
+We'll update this implementation with the new Triton backend once this is fixed.
+
+We use the FlashAttention implementation from Phil Tillet a starting point.
+https://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py
+
+Changes:
+- Implement both causal and non-causal attention.
+- Implement both self-attention and cross-attention.
+- Support arbitrary seqlens (not just multiples of 128), for both forward and backward.
+- Support all head dimensions up to 128 (not just 16, 32, 64, 128), for both forward and backward.
+- Support attention bias.
+- Speed up the forward pass a bit, and only store the LSE instead of m and l.
+- Make the backward for d=128 much faster by reducing register spilling.
+- Optionally parallelize the backward pass across seqlen_k, to deal with the case of
+small batch size * nheads.
+
+Caution:
+- This is an *experimental* implementation. The forward pass should be quite robust but
+I'm not 100% sure that the backward pass doesn't have race conditions (due to the Triton compiler).
+- This implementation has only been tested on A100.
+- If you plan to use headdim other than 64 and 128, you should test for race conditions
+(due to the Triton compiler), as done in tests/test_flash_attn.py
+"test_flash_attn_triton_race_condition". I've tested and fixed many race conditions
+for different head dimensions (40, 48, 64, 128, 80, 88, 96), but I'm still not 100% confident
+that there are none left for other head dimensions.
+
+Differences between this Triton version and the CUDA version:
+- Triton version doesn't support dropout.
+- Triton forward is generally faster than CUDA forward, while Triton backward is
+generally slower than CUDA backward. Overall Triton forward + backward is slightly slower
+than CUDA forward + backward.
+- Triton version doesn't support different sequence lengths in a batch (i.e., RaggedTensor/NestedTensor).
+- Triton version supports attention bias, while CUDA version doesn't.
+"""
+
+import math
+
+import torch
+import os
+
+import triton_pre_mlir as triton
+import triton_pre_mlir.compiler
+import triton_pre_mlir.language as tl
+import functools
+import subprocess
+
+if 'CONDA_PREFIX' in os.environ and 'CUDA_HOME' not in os.environ:
+    os.environ['CUDA_HOME'] = os.environ['CONDA_PREFIX']
+
+
+@functools.lru_cache()
+def libcuda_dirs():
+    libs = subprocess.check_output(["ldconfig", "-p"]).decode()
+    # each line looks like the following:
+    # libcuda.so.1 (libc6,x86-64) => /lib/x86_64-linux-gnu/libcuda.so.1
+    locs = [line.split()[-1] for line in libs.splitlines() if "libcuda.so" in line]
+    dirs = [os.path.dirname(loc) for loc in locs]
+    msg = 'libcuda.so cannot found!\n'
+    if locs:
+        msg += 'Possible files are located at %s.' % str(locs)
+        msg += 'Please create a symlink of libcuda.so to any of the file.'
+    assert any(os.path.exists(os.path.join(path, 'libcuda.so')) for path in dirs), msg
+    return dirs
+
+
+triton_pre_mlir.compiler.libcuda_dirs = libcuda_dirs
+
+
+# Disabling autotune for now, set num_warps=4 if headdim=64 and num_warps=8 if headdim=128
+# @triton.autotune(
+#     configs=[
+#         triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=4, num_stages=1),
+#         # This config has a race condition when EVEN_M == False, disabling it for now.
+#         # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_warps=4, num_stages=1),
+#     ],
+#     key=['CACHE_KEY_SEQLEN_Q', 'CACHE_KEY_SEQLEN_K', 'BIAS_TYPE', 'IS_CAUSAL', 'BLOCK_HEADDIM']
+# )
+@triton.heuristics(
+    {
+        "EVEN_M": lambda args: args["seqlen_q"] % args["BLOCK_M"] == 0,
+        "EVEN_N": lambda args: args["seqlen_k"] % args["BLOCK_N"] == 0,
+        "EVEN_HEADDIM": lambda args: args["headdim"] == args["BLOCK_HEADDIM"],
+    }
+)
+@triton.jit
+def _fwd_kernel(
+        Q, K, V, Bias, Out,
+        Lse, TMP,  # NOTE: TMP is a scratchpad buffer to workaround a compiler bug
+        softmax_scale,
+        stride_qb, stride_qh, stride_qm,
+        stride_kb, stride_kh, stride_kn,
+        stride_vb, stride_vh, stride_vn,
+        stride_bb, stride_bh, stride_bm,
+        stride_ob, stride_oh, stride_om,
+        nheads, seqlen_q, seqlen_k, seqlen_q_rounded, headdim,
+        CACHE_KEY_SEQLEN_Q, CACHE_KEY_SEQLEN_K,
+        BIAS_TYPE: tl.constexpr,
+        IS_CAUSAL: tl.constexpr,
+        BLOCK_HEADDIM: tl.constexpr,
+        EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr,
+        BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
+):
+    start_m = tl.program_id(0)
+    off_hb = tl.program_id(1)
+    off_b = off_hb // nheads
+    off_h = off_hb % nheads
+    # off_b = tl.program_id(1)
+    # off_h = tl.program_id(2)
+    # off_hb = off_b * nheads + off_h
+    # initialize offsets
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, BLOCK_HEADDIM)
+    # Initialize pointers to Q, K, V
+    # Adding parenthesis around indexing might use int32 math instead of int64 math?
+    # https://github.com/openai/triton/issues/741
+    # I'm seeing a tiny bit of difference (5-7us)
+    q_ptrs = Q + off_b * stride_qb + off_h * stride_qh + (offs_m[:, None] * stride_qm + offs_d[None, :])
+    k_ptrs = K + off_b * stride_kb + off_h * stride_kh + (offs_n[:, None] * stride_kn + offs_d[None, :])
+    v_ptrs = V + off_b * stride_vb + off_h * stride_vh + (offs_n[:, None] * stride_vn + offs_d[None, :])
+    if BIAS_TYPE == 'vector':
+        b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + offs_n
+    elif BIAS_TYPE == 'matrix':
+        b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + (offs_m[:, None] * stride_bm + offs_n[None, :])
+    # initialize pointer to m and l
+    t_ptrs = TMP + off_hb * seqlen_q_rounded + offs_m
+    lse_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+    acc_o = tl.zeros([BLOCK_M, BLOCK_HEADDIM], dtype=tl.float32)
+    # load q: it will stay in SRAM throughout
+    # [2022-10-30] TD: Triton bug - in the case of EVEN_M=True and EVEN_N=False, if we just call
+    # tl.load(q_ptrs), we get the wrong output!
+    if EVEN_M & EVEN_N:
+        if EVEN_HEADDIM:
+            q = tl.load(q_ptrs)
+        else:
+            q = tl.load(q_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
+    else:
+        if EVEN_HEADDIM:
+            q = tl.load(q_ptrs, mask=offs_m[:, None] < seqlen_q, other=0.0)
+        else:
+            q = tl.load(q_ptrs, mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim),
+                        other=0.0)
+    # loop over k, v and update accumulator
+    end_n = seqlen_k if not IS_CAUSAL else tl.minimum((start_m + 1) * BLOCK_M, seqlen_k)
+    for start_n in range(0, end_n, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        # -- compute qk ----
+        if EVEN_N & EVEN_M:  # If we just do "if EVEN_N", there seems to be some race condition
+            if EVEN_HEADDIM:
+                k = tl.load(k_ptrs + start_n * stride_kn)
+            else:
+                k = tl.load(k_ptrs + start_n * stride_kn, mask=offs_d[None, :] < headdim, other=0.0)
+        else:
+            if EVEN_HEADDIM:
+                k = tl.load(k_ptrs + start_n * stride_kn, mask=(start_n + offs_n)[:, None] < seqlen_k,
+                            other=0.0)
+            else:
+                k = tl.load(k_ptrs + start_n * stride_kn,
+                            mask=((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim),
+                            other=0.0)
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        qk += tl.dot(q, k, trans_b=True)
+        # Trying to combine the two masks seem to make the result wrong
+        if not EVEN_N:  # Need to mask out otherwise the softmax is wrong
+            qk += tl.where((start_n + offs_n)[None, :] < seqlen_k, 0, float("-inf"))
+        if IS_CAUSAL:
+            qk += tl.where(offs_m[:, None] >= (start_n + offs_n)[None, :], 0, float("-inf"))
+        if BIAS_TYPE != 'none':
+            if BIAS_TYPE == 'vector':
+                if EVEN_N:
+                    bias = tl.load(b_ptrs + start_n).to(tl.float32)
+                else:
+                    bias = tl.load(b_ptrs + start_n, mask=(start_n + offs_n) < seqlen_k, other=0.0).to(tl.float32)
+                bias = bias[None, :]
+            elif BIAS_TYPE == 'matrix':
+                if EVEN_M & EVEN_N:
+                    bias = tl.load(b_ptrs + start_n).to(tl.float32)
+                else:
+                    bias = tl.load(b_ptrs + start_n,
+                                   mask=(offs_m[:, None] < seqlen_q)
+                                        & ((start_n + offs_n)[None, :] < seqlen_k),
+                                   other=0.0).to(tl.float32)
+            # Slightly faster to multiply the softmax_scale in the tl.exp below since the compiler
+            # can then fuse the mult and add into an fma instruction. But if we have bias we need to
+            # to multiply with softmax_scale here.
+            qk = qk * softmax_scale + bias
+            m_ij = tl.maximum(tl.max(qk, 1), lse_i)
+            p = tl.exp(qk - m_ij[:, None])
+        else:
+            m_ij = tl.maximum(tl.max(qk, 1) * softmax_scale, lse_i)
+            p = tl.exp(qk * softmax_scale - m_ij[:, None])
+        l_ij = tl.sum(p, 1)
+
+        # scale acc_o
+        acc_o_scale = tl.exp(m_i - m_ij)
+
+        # # -- update output accumulator --
+        # BUG: have to store and immediately load
+        tl.store(t_ptrs, acc_o_scale)
+        acc_o_scale = tl.load(t_ptrs)
+        acc_o = acc_o * acc_o_scale[:, None]
+        # update acc_o
+        if EVEN_N & EVEN_M:  # If we just do "if EVEN_N", there seems to be some race condition
+            if EVEN_HEADDIM:
+                v = tl.load(v_ptrs + start_n * stride_vn)
+            else:
+                v = tl.load(v_ptrs + start_n * stride_vn, mask=offs_d[None, :] < headdim, other=0.0)
+        else:
+            if EVEN_HEADDIM:
+                v = tl.load(v_ptrs + start_n * stride_vn, mask=(start_n + offs_n)[:, None] < seqlen_k,
+                            other=0.0)
+            else:
+                v = tl.load(v_ptrs + start_n * stride_vn,
+                            mask=((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim),
+                            other=0.0)
+        p = p.to(v.dtype)
+        acc_o += tl.dot(p, v)
+
+        # -- update statistics
+        m_i = m_ij
+        l_i_new = tl.exp(lse_i - m_ij) + l_ij
+        lse_i = m_ij + tl.log(l_i_new)
+
+    o_scale = tl.exp(m_i - lse_i)
+    # BUG: have to store and immediately load
+    tl.store(t_ptrs, o_scale)
+    o_scale = tl.load(t_ptrs)
+    acc_o = acc_o * o_scale[:, None]
+    # rematerialize offsets to save registers
+    start_m = tl.program_id(0)
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    # write back l and m
+    lse_ptrs = Lse + off_hb * seqlen_q_rounded + offs_m
+    tl.store(lse_ptrs, lse_i)
+    # initialize pointers to output
+    offs_d = tl.arange(0, BLOCK_HEADDIM)
+    out_ptrs = Out + off_b * stride_ob + off_h * stride_oh + (offs_m[:, None] * stride_om + offs_d[None, :])
+    if EVEN_M:
+        if EVEN_HEADDIM:
+            tl.store(out_ptrs, acc_o)
+        else:
+            tl.store(out_ptrs, acc_o, mask=offs_d[None, :] < headdim)
+    else:
+        if EVEN_HEADDIM:
+            tl.store(out_ptrs, acc_o, mask=offs_m[:, None] < seqlen_q)
+        else:
+            tl.store(out_ptrs, acc_o,
+                     mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim))
+
+
+@triton.jit
+def _bwd_preprocess_do_o_dot(
+        Out, DO, Delta,
+        stride_ob, stride_oh, stride_om,
+        stride_dob, stride_doh, stride_dom,
+        nheads, seqlen_q, seqlen_q_rounded, headdim,
+        BLOCK_M: tl.constexpr, BLOCK_HEADDIM: tl.constexpr,
+):
+    start_m = tl.program_id(0)
+    off_hb = tl.program_id(1)
+    off_b = off_hb // nheads
+    off_h = off_hb % nheads
+    # initialize offsets
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, BLOCK_HEADDIM)
+    # load
+    o = tl.load(Out + off_b * stride_ob + off_h * stride_oh + offs_m[:, None] * stride_om + offs_d[None, :],
+                mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0).to(tl.float32)
+    do = tl.load(DO + off_b * stride_dob + off_h * stride_doh + offs_m[:, None] * stride_dom + offs_d[None, :],
+                 mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0).to(tl.float32)
+    delta = tl.sum(o * do, axis=1)
+    # write-back
+    tl.store(Delta + off_hb * seqlen_q_rounded + offs_m, delta)
+
+
+@triton.jit
+def _bwd_store_dk_dv(
+        dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim,
+        EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr,
+):
+    # [2022-11-01] TD: Same bug. In the case of EVEN_N=True and EVEN_M=False,
+    # if we just call tl.store(dv_ptrs), there's a race condition
+    if EVEN_N & EVEN_M:
+        if EVEN_HEADDIM:
+            tl.store(dv_ptrs, dv)
+            tl.store(dk_ptrs, dk)
+        else:
+            tl.store(dv_ptrs, dv, mask=offs_d[None, :] < headdim)
+            tl.store(dk_ptrs, dk, mask=offs_d[None, :] < headdim)
+    else:
+        if EVEN_HEADDIM:
+            tl.store(dv_ptrs, dv, mask=offs_n[:, None] < seqlen_k)
+            tl.store(dk_ptrs, dk, mask=offs_n[:, None] < seqlen_k)
+        else:
+            tl.store(dv_ptrs, dv, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim))
+            tl.store(dk_ptrs, dk, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim))
+
+
+@triton.jit
+def _bwd_kernel_one_col_block(
+        start_n,
+        Q, K, V, Bias,
+        DO, DQ, DK, DV,
+        LSE, D,
+        softmax_scale,
+        stride_qm, stride_kn, stride_vn, stride_bm,
+        stride_dom, stride_dqm, stride_dkn, stride_dvn,
+        seqlen_q, seqlen_k, headdim,
+        ATOMIC_ADD: tl.constexpr,
+        BIAS_TYPE: tl.constexpr,
+        IS_CAUSAL: tl.constexpr,
+        BLOCK_HEADDIM: tl.constexpr,
+        EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr,
+        BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
+):
+    # We need to make sure begin_m is a multiple of BLOCK_M (not BLOCK_N)
+    begin_m = 0 if not IS_CAUSAL else ((start_n * BLOCK_N) // BLOCK_M) * BLOCK_M
+    # initialize row/col offsets
+    offs_qm = begin_m + tl.arange(0, BLOCK_M)
+    offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_m = tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, BLOCK_HEADDIM)
+    # initialize pointers to value-like data
+    q_ptrs = Q + (offs_qm[:, None] * stride_qm + offs_d[None, :])
+    k_ptrs = K + (offs_n[:, None] * stride_kn + offs_d[None, :])
+    v_ptrs = V + (offs_n[:, None] * stride_vn + offs_d[None, :])
+    do_ptrs = DO + (offs_qm[:, None] * stride_dom + offs_d[None, :])
+    dq_ptrs = DQ + (offs_qm[:, None] * stride_dqm + offs_d[None, :])
+    if BIAS_TYPE == 'vector':
+        b_ptrs = Bias + offs_n
+    elif BIAS_TYPE == 'matrix':
+        b_ptrs = Bias + (offs_qm[:, None] * stride_bm + offs_n[None, :])
+    # initialize dv and dk
+    dv = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
+    dk = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
+    # There seems to be some problem with Triton pipelining that makes results wrong for
+    # headdim=64, seqlen=(113, 255), bias_type='matrix'. In this case the for loop
+    # may have zero step, and pipelining with the bias matrix could screw it up.
+    # So we just exit early.
+    if begin_m >= seqlen_q:
+        dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :])
+        dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :])
+        _bwd_store_dk_dv(dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim,
+                         EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM)
+        return
+    # k and v stay in SRAM throughout
+    # [2022-10-30] TD: Same bug as the fwd. In the case of EVEN_N=True and EVEN_M=False,
+    # if we just call tl.load(k_ptrs), we get the wrong output!
+    if EVEN_N & EVEN_M:
+        if EVEN_HEADDIM:
+            k = tl.load(k_ptrs)
+            v = tl.load(v_ptrs)
+        else:
+            k = tl.load(k_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
+            v = tl.load(v_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
+    else:
+        if EVEN_HEADDIM:
+            k = tl.load(k_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)
+            v = tl.load(v_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)
+        else:
+            k = tl.load(k_ptrs, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim),
+                        other=0.0)
+            v = tl.load(v_ptrs, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim),
+                        other=0.0)
+    # loop over rows
+    num_block_m = tl.cdiv(seqlen_q, BLOCK_M)
+    for start_m in range(begin_m, num_block_m * BLOCK_M, BLOCK_M):
+        start_m = tl.multiple_of(start_m, BLOCK_M)
+        offs_m_curr = start_m + offs_m
+        # load q, k, v, do on-chip
+        # Same bug as below. Otherwise gives wrong result for headdim=40, seqlen=(128, 117)
+        if EVEN_M & EVEN_HEADDIM:
+            q = tl.load(q_ptrs)
+        else:
+            if EVEN_HEADDIM:
+                q = tl.load(q_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0)
+            else:
+                q = tl.load(q_ptrs, mask=(offs_m_curr[:, None] < seqlen_q)
+                                         & (offs_d[None, :] < headdim), other=0.0)
+        # recompute p = softmax(qk, dim=-1).T
+        qk = tl.dot(q, k, trans_b=True)
+        # Trying to combine the two masks seem to make the result wrong
+        if not EVEN_N:  # Need to mask out otherwise the softmax is wrong
+            qk = tl.where(offs_n[None, :] < seqlen_k, qk, float("-inf"))
+        if IS_CAUSAL:
+            qk = tl.where(offs_m_curr[:, None] >= (offs_n[None, :]), qk, float("-inf"))
+        if BIAS_TYPE != 'none':
+            tl.debug_barrier()  # Race condition otherwise
+            if BIAS_TYPE == 'vector':
+                if EVEN_N:
+                    bias = tl.load(b_ptrs).to(tl.float32)
+                else:
+                    bias = tl.load(b_ptrs, mask=offs_n < seqlen_k, other=0.0).to(tl.float32)
+                bias = bias[None, :]
+            elif BIAS_TYPE == 'matrix':
+                if EVEN_M & EVEN_N:
+                    bias = tl.load(b_ptrs).to(tl.float32)
+                else:
+                    bias = tl.load(b_ptrs,
+                                   mask=(offs_m_curr[:, None] < seqlen_q)
+                                        & (offs_n[None, :] < seqlen_k),
+                                   other=0.0).to(tl.float32)
+            qk = qk * softmax_scale + bias
+        # There seems to be a race condition when headdim=48/96, and dq, dk, dv are wrong.
+        # Also wrong for headdim=64.
+        if not (EVEN_M & EVEN_HEADDIM):
+            tl.debug_barrier()
+        lse_i = tl.load(LSE + offs_m_curr)
+        if BIAS_TYPE == 'none':
+            p = tl.exp(qk * softmax_scale - lse_i[:, None])
+        else:
+            p = tl.exp(qk - lse_i[:, None])
+        # compute dv
+        # [2022-10-30] TD: A Triton bug: if EVEN_M=True and EVEN_HEADDIM=False, if we call
+        # do = tl.load(do_ptrs, mask=offs_d[None, :] < headdim, other=0.0), we get wrong outputs
+        # in the case of headdim=48/96, seqlen_q & seqlen_k >= 512. If headdim=40 or seqlen < 512,
+        # the output is correct.
+        if EVEN_M & EVEN_HEADDIM:
+            do = tl.load(do_ptrs)
+        else:
+            # [2022-11-01] TD: Triton bug, there's a race condition if we just use m_mask and not d_mask.
+            do = tl.load(do_ptrs, mask=(offs_m_curr[:, None] < seqlen_q)
+                                       & (offs_d[None, :] < headdim), other=0.0)
+        # if EVEN_M:
+        #     if EVEN_HEADDIM:
+        #         do = tl.load(do_ptrs)
+        #     else:
+        #         do = tl.load(do_ptrs, mask=offs_d[None, :] < headdim, other=0.0)
+        # else:
+        #     if EVEN_HEADDIM:
+        #         do = tl.load(do_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0)
+        #     else:
+        #         do = tl.load(do_ptrs, mask=(offs_m_curr[:, None] < seqlen_q)
+        #                                    & (offs_d[None, :] < headdim), other=0.0)
+        dv += tl.dot(p.to(do.dtype), do, trans_a=True)
+        # compute dp = dot(v, do)
+        # There seems to be a race condition when headdim=48/96, and dq, dk are wrong.
+        # Also wrong for headdim=128, seqlen=(108, 256), and ATOMIC_ADD=True
+        # Also wrong for headdim=64, seqlen=(1023, 1024), and ATOMIC_ADD=False
+        if not (EVEN_M & EVEN_HEADDIM):
+            tl.debug_barrier()
+        dp = tl.dot(do, v, trans_b=True)
+        # There's a race condition for headdim=48
+        if not EVEN_HEADDIM:
+            tl.debug_barrier()
+        # compute ds = p * (dp - delta[:, None])
+        # Putting the subtraction after the dp matmul (instead of before) is slightly faster
+        Di = tl.load(D + offs_m_curr)
+        # Converting ds to q.dtype here reduces register pressure and makes it much faster
+        # for BLOCK_HEADDIM=128
+        ds = (p * (dp - Di[:, None]) * softmax_scale).to(q.dtype)
+        # compute dk = dot(ds.T, q)
+        dk += tl.dot(ds, q, trans_a=True)
+        # compute dq
+        if not (EVEN_M & EVEN_HEADDIM):  # Otherewise there's a race condition when BIAS_TYPE='matrix'
+            tl.debug_barrier()
+        if not ATOMIC_ADD:
+            if EVEN_M & EVEN_HEADDIM:  # Race condition if we just do EVEN_M
+                dq = tl.load(dq_ptrs, eviction_policy="evict_last")
+                dq += tl.dot(ds, k)
+                tl.store(dq_ptrs, dq, eviction_policy="evict_last")
+            else:
+                if EVEN_HEADDIM:
+                    dq = tl.load(dq_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0,
+                                 eviction_policy="evict_last")
+                    dq += tl.dot(ds, k)
+                    tl.store(dq_ptrs, dq, mask=offs_m_curr[:, None] < seqlen_q,
+                             eviction_policy="evict_last")
+                else:
+                    dq = tl.load(dq_ptrs,
+                                 mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim),
+                                 other=0.0, eviction_policy="evict_last")
+                    dq += tl.dot(ds, k)
+                    tl.store(dq_ptrs, dq,
+                             mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim),
+                             eviction_policy="evict_last")
+        else:  # If we're parallelizing across the seqlen_k dimension
+            dq = tl.dot(ds, k)
+            if EVEN_M & EVEN_HEADDIM:  # Race condition if we just do EVEN_M
+                tl.atomic_add(dq_ptrs, dq)
+            else:
+                if EVEN_HEADDIM:
+                    tl.atomic_add(dq_ptrs, dq, mask=offs_m_curr[:, None] < seqlen_q)
+                else:
+                    tl.atomic_add(dq_ptrs, dq,
+                                  mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim))
+        # increment pointers
+        dq_ptrs += BLOCK_M * stride_dqm
+        q_ptrs += BLOCK_M * stride_qm
+        do_ptrs += BLOCK_M * stride_dom
+        if BIAS_TYPE == 'matrix':
+            b_ptrs += BLOCK_M * stride_bm
+    # write-back
+    dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :])
+    dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :])
+    _bwd_store_dk_dv(dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim,
+                     EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM)
+
+
+def init_to_zero(name):
+    return lambda nargs: nargs[name].zero_()
+
+
+# TODO: Change BLOCK_M and BLOCK_N according to your GPU and num_warps according to headdim
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "SEQUENCE_PARALLEL": False}, num_warps=4, num_stages=1, pre_hook=init_to_zero('DQ')),
+        triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "SEQUENCE_PARALLEL": True}, num_warps=4, num_stages=1, pre_hook=init_to_zero('DQ')),
+        # Other configs seem to give wrong results when seqlen_q % 128 != 0, disabling them for now
+        # # Kernel is buggy (give wrong result) if we set BLOCK_m=128, BLOCK_n=64, num_warps=*4*
+        # triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "SEQUENCE_PARALLEL": False}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ')),
+        # triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "SEQUENCE_PARALLEL": True}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ')),
+        # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "SEQUENCE_PARALLEL": False}, num_warps=4, num_stages=1, pre_hook=init_to_zero('DQ')),
+        # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "SEQUENCE_PARALLEL": True}, num_warps=4, num_stages=1, pre_hook=init_to_zero('DQ')),
+    ],
+    key=['CACHE_KEY_SEQLEN_Q', 'CACHE_KEY_SEQLEN_K', 'BIAS_TYPE', 'IS_CAUSAL', 'BLOCK_HEADDIM'],
+)
+@triton.heuristics(
+    {
+        "EVEN_M": lambda args: args["seqlen_q"] % args["BLOCK_M"] == 0,
+        "EVEN_N": lambda args: args["seqlen_k"] % args["BLOCK_N"] == 0,
+        "EVEN_HEADDIM": lambda args: args["headdim"] == args["BLOCK_HEADDIM"],
+    }
+)
+@triton.jit
+def _bwd_kernel(
+        Q, K, V, Bias,
+        DO, DQ, DK, DV,
+        LSE, D,
+        softmax_scale,
+        stride_qb, stride_qh, stride_qm,
+        stride_kb, stride_kh, stride_kn,
+        stride_vb, stride_vh, stride_vn,
+        stride_bb, stride_bh, stride_bm,
+        stride_dob, stride_doh, stride_dom,
+        stride_dqb, stride_dqh, stride_dqm,
+        stride_dkb, stride_dkh, stride_dkn,
+        stride_dvb, stride_dvh, stride_dvn,
+        nheads, seqlen_q, seqlen_k, seqlen_q_rounded, headdim,
+        CACHE_KEY_SEQLEN_Q, CACHE_KEY_SEQLEN_K,
+        BIAS_TYPE: tl.constexpr,
+        IS_CAUSAL: tl.constexpr,
+        BLOCK_HEADDIM: tl.constexpr,
+        SEQUENCE_PARALLEL: tl.constexpr,
+        EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr,
+        BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
+):
+    off_hb = tl.program_id(1)
+    off_b = off_hb // nheads
+    off_h = off_hb % nheads
+    # offset pointers for batch/head
+    Q += off_b * stride_qb + off_h * stride_qh
+    K += off_b * stride_kb + off_h * stride_kh
+    V += off_b * stride_vb + off_h * stride_vh
+    DO += off_b * stride_dob + off_h * stride_doh
+    DQ += off_b * stride_dqb + off_h * stride_dqh
+    DK += off_b * stride_dkb + off_h * stride_dkh
+    DV += off_b * stride_dvb + off_h * stride_dvh
+    if BIAS_TYPE != 'none':
+        Bias += off_b * stride_bb + off_h * stride_bh
+    # pointer to row-wise quantities in value-like data
+    D += off_hb * seqlen_q_rounded
+    LSE += off_hb * seqlen_q_rounded
+    if not SEQUENCE_PARALLEL:
+        num_block_n = tl.cdiv(seqlen_k, BLOCK_N)
+        for start_n in range(0, num_block_n):
+            _bwd_kernel_one_col_block(
+                start_n,
+                Q, K, V, Bias,
+                DO, DQ, DK, DV,
+                LSE, D,
+                softmax_scale,
+                stride_qm, stride_kn, stride_vn, stride_bm,
+                stride_dom, stride_dqm, stride_dkn, stride_dvn,
+                seqlen_q, seqlen_k, headdim,
+                ATOMIC_ADD=False,
+                BIAS_TYPE=BIAS_TYPE,
+                IS_CAUSAL=IS_CAUSAL,
+                BLOCK_HEADDIM=BLOCK_HEADDIM,
+                EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM,
+                BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N
+            )
+    else:
+        start_n = tl.program_id(0)
+        _bwd_kernel_one_col_block(
+            start_n,
+            Q, K, V, Bias,
+            DO, DQ, DK, DV,
+            LSE, D,
+            softmax_scale,
+            stride_qm, stride_kn, stride_vn, stride_bm,
+            stride_dom, stride_dqm, stride_dkn, stride_dvn,
+            seqlen_q, seqlen_k, headdim,
+            ATOMIC_ADD=True,
+            BIAS_TYPE=BIAS_TYPE,
+            IS_CAUSAL=IS_CAUSAL,
+            BLOCK_HEADDIM=BLOCK_HEADDIM,
+            EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM,
+            BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N
+        )
+
+
+def _flash_attn_forward(q, k, v, bias=None, causal=False, softmax_scale=None):
+    # shape constraints
+    batch, seqlen_q, nheads, d = q.shape
+    _, seqlen_k, _, _ = k.shape
+    assert k.shape == (batch, seqlen_k, nheads, d)
+    assert v.shape == (batch, seqlen_k, nheads, d)
+    assert d <= 128, 'FlashAttention only support head dimensions up to 128'
+    assert q.dtype == k.dtype == v.dtype, 'All tensors must have the same type'
+    assert q.dtype in [torch.float16, torch.bfloat16], 'Only support fp16 and bf16'
+    assert q.is_cuda and k.is_cuda and v.is_cuda
+    softmax_scale = softmax_scale or 1.0 / math.sqrt(d)
+
+    has_bias = bias is not None
+    bias_type = 'none'
+    if has_bias:
+        assert bias.dtype in [q.dtype, torch.float]
+        assert bias.is_cuda
+        assert bias.dim() == 4
+        if bias.stride(-1) != 1:
+            bias = bias.contiguous()
+        if bias.shape[2:] == (1, seqlen_k):
+            bias_type = 'vector'
+        elif bias.shape[2:] == (seqlen_q, seqlen_k):
+            bias_type = 'matrix'
+        else:
+            raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k)'
+                               ' or (seqlen_q, seqlen_k)')
+        bias = bias.expand(batch, nheads, seqlen_q, seqlen_k)
+    bias_strides = (bias.stride(0), bias.stride(1), bias.stride(2)) if has_bias else (0, 0, 0)
+
+    seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128
+    lse = torch.empty((batch, nheads, seqlen_q_rounded), device=q.device, dtype=torch.float32)
+    tmp = torch.empty((batch, nheads, seqlen_q_rounded), device=q.device, dtype=torch.float32)
+    o = torch.empty_like(q)
+
+    BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)
+    BLOCK = 128
+    num_warps = 4 if d <= 64 else 8
+    grid = lambda META: (triton.cdiv(seqlen_q, META["BLOCK_M"]), batch * nheads)
+    _fwd_kernel[grid](
+        q, k, v, bias, o,
+        lse, tmp,
+        softmax_scale,
+        q.stride(0), q.stride(2), q.stride(1),
+        k.stride(0), k.stride(2), k.stride(1),
+        v.stride(0), v.stride(2), v.stride(1),
+        *bias_strides,
+        o.stride(0), o.stride(2), o.stride(1),
+        nheads, seqlen_q, seqlen_k, seqlen_q_rounded, d,
+        seqlen_q // 32, seqlen_k // 32,  # key for triton cache (limit number of compilations)
+        # Can't use kwargs here because triton autotune expects key to be args, not kwargs
+        # IS_CAUSAL=causal, BLOCK_HEADDIM=d,
+        bias_type, causal, BLOCK_HEADDIM,
+        BLOCK_M=BLOCK, BLOCK_N=BLOCK,
+        num_warps=num_warps,
+        num_stages=1,
+    )
+    return o, lse, softmax_scale  # softmax_scale could have been updated
+
+
+def _flash_attn_backward(do, q, k, v, o, lse, dq, dk, dv, bias=None, causal=False, softmax_scale=None):
+    # Make sure that the last dimension is contiguous
+    if do.stride(-1) != 1:
+        do = do.contiguous()
+    batch, seqlen_q, nheads, d = q.shape
+    _, seqlen_k, _, _ = k.shape
+    # assert d in {16, 32, 64, 128}
+    assert d <= 128
+    seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128
+    assert lse.shape == (batch, nheads, seqlen_q_rounded)
+    assert q.stride(-1) == k.stride(-1) == v.stride(-1) == o.stride(-1) == 1
+    assert dq.stride(-1) == dk.stride(-1) == dv.stride(-1) == 1
+    softmax_scale = softmax_scale or 1.0 / math.sqrt(d)
+    # dq_accum = torch.zeros_like(q, dtype=torch.float32)
+    dq_accum = torch.empty_like(q, dtype=torch.float32)
+    delta = torch.empty_like(lse)
+    # delta = torch.zeros_like(lse)
+
+    BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)
+    grid = lambda META: (triton.cdiv(seqlen_q, META["BLOCK_M"]), batch * nheads)
+    _bwd_preprocess_do_o_dot[grid](
+        o, do, delta,
+        o.stride(0), o.stride(2), o.stride(1),
+        do.stride(0), do.stride(2), do.stride(1),
+        nheads, seqlen_q, seqlen_q_rounded, d,
+        BLOCK_M=128, BLOCK_HEADDIM=BLOCK_HEADDIM,
+    )
+
+    has_bias = bias is not None
+    bias_type = 'none'
+    if has_bias:
+        assert bias.dtype in [q.dtype, torch.float]
+        assert bias.is_cuda
+        assert bias.dim() == 4
+        assert bias.stride(-1) == 1
+        if bias.shape[2:] == (1, seqlen_k):
+            bias_type = 'vector'
+        elif bias.shape[2:] == (seqlen_q, seqlen_k):
+            bias_type = 'matrix'
+        else:
+            raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k)'
+                               ' or (seqlen_q, seqlen_k)')
+        bias = bias.expand(batch, nheads, seqlen_q, seqlen_k)
+    bias_strides = (bias.stride(0), bias.stride(1), bias.stride(2)) if has_bias else (0, 0, 0)
+
+    # BLOCK_M = 128
+    # BLOCK_N = 64
+    # num_warps = 4
+    grid = lambda META: (triton.cdiv(seqlen_k, META["BLOCK_N"]) if META["SEQUENCE_PARALLEL"] else 1,
+                         batch * nheads)
+    _bwd_kernel[grid](
+        q, k, v, bias,
+        do, dq_accum, dk, dv,
+        lse, delta,
+        softmax_scale,
+        q.stride(0), q.stride(2), q.stride(1),
+        k.stride(0), k.stride(2), k.stride(1),
+        v.stride(0), v.stride(2), v.stride(1),
+        *bias_strides,
+        do.stride(0), do.stride(2), do.stride(1),
+        dq_accum.stride(0), dq_accum.stride(2), dq_accum.stride(1),
+        dk.stride(0), dk.stride(2), dk.stride(1),
+        dv.stride(0), dv.stride(2), dv.stride(1),
+        nheads, seqlen_q, seqlen_k, seqlen_q_rounded, d,
+        seqlen_q // 32, seqlen_k // 32,  # key for triton cache (limit number of compilations)
+        # Can't use kwargs here because triton autotune expects key to be args, not kwargs
+        # IS_CAUSAL=causal, BLOCK_HEADDIM=d,
+        bias_type, causal, BLOCK_HEADDIM,
+        # SEQUENCE_PARALLEL=False,
+        # BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
+        # num_warps=num_warps,
+        # num_stages=1,
+    )
+    dq.copy_(dq_accum)
+
+
+class FlashAttnQKVPackedFunc(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, qkv, bias=None, causal=False, softmax_scale=None):
+        """
+            qkv: (batch, seqlen, 3, nheads, headdim)
+            bias: optional, shape broadcastible to (batch, nheads, seqlen, seqlen).
+                For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen).
+                ALiBi mask for non-causal would have shape (1, nheads, seqlen, seqlen)
+        """
+        # Make sure that the last dimension is contiguous
+        if qkv.stride(-1) != 1:
+            qkv = qkv.contiguous()
+        o, lse, ctx.softmax_scale = _flash_attn_forward(
+            qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2], bias=bias, causal=causal,
+            softmax_scale=softmax_scale
+        )
+        ctx.save_for_backward(qkv, o, lse, bias)
+        ctx.causal = causal
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        qkv, o, lse, bias = ctx.saved_tensors
+        assert not ctx.needs_input_grad[1], 'FlashAttention does not support bias gradient yet'
+        # Triton's autotune causes the Tensor._version to change, and so Pytorch autograd
+        # does a memcpy. To avoid this we run in inference_mode, which doesn't track the version.
+        with torch.inference_mode():
+            dqkv = torch.empty_like(qkv)
+            _flash_attn_backward(do, qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2], o, lse,
+                                 dqkv[:, :, 0], dqkv[:, :, 1], dqkv[:, :, 2],
+                                 bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale)
+        return dqkv, None, None, None
+
+
+flash_attn_qkvpacked_func = FlashAttnQKVPackedFunc.apply
+
+
+class FlashAttnKVPackedFunc(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, q, kv, bias=None, causal=False, softmax_scale=None):
+        """
+            q: (batch, seqlen_q, nheads, headdim)
+            kv: (batch, seqlen_k, 2, nheads, headdim)
+            bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).
+                For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).
+                ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)
+        """
+        # Make sure that the last dimension is contiguous
+        q, kv = [x if x.stride(-1) == 1 else x.contiguous() for x in [q, kv]]
+        o, lse, ctx.softmax_scale = _flash_attn_forward(
+            q, kv[:, :, 0], kv[:, :, 1], bias=bias, causal=causal, softmax_scale=softmax_scale
+        )
+        ctx.save_for_backward(q, kv, o, lse, bias)
+        ctx.causal = causal
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        q, kv, o, lse, bias = ctx.saved_tensors
+        if len(ctx.needs_input_grad) >= 3:
+            assert not ctx.needs_input_grad[2], 'FlashAttention does not support bias gradient yet'
+        # Triton's autotune causes the Tensor._version to change, and so Pytorch autograd
+        # does a memcpy. To avoid this we run in inference_mode, which doesn't track the version.
+        with torch.inference_mode():
+            dq = torch.empty_like(q)
+            dkv = torch.empty_like(kv)
+            _flash_attn_backward(do, q, kv[:, :, 0], kv[:, :, 1], o, lse,
+                                 dq, dkv[:, :, 0], dkv[:, :, 1],
+                                 bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale)
+        return dq, dkv, None, None, None
+
+
+flash_attn_kvpacked_func = FlashAttnKVPackedFunc.apply
+
+
+class FlashAttnFunc(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, q, k, v, bias=None, causal=False, softmax_scale=None):
+        """
+            q: (batch_size, seqlen_q, nheads, headdim)
+            k, v: (batch_size, seqlen_k, nheads, headdim)
+            bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).
+                For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).
+                ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)
+        """
+        # Make sure that the last dimension is contiguous
+        q, k, v = [x if x.stride(-1) == 1 else x.contiguous() for x in [q, k, v]]
+        o, lse, ctx.softmax_scale = _flash_attn_forward(
+            q, k, v, bias=bias, causal=causal, softmax_scale=softmax_scale
+        )
+        ctx.save_for_backward(q, k, v, o, lse, bias)
+        ctx.causal = causal
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        q, k, v, o, lse, bias = ctx.saved_tensors
+        assert not ctx.needs_input_grad[3], 'FlashAttention does not support bias gradient yet'
+        # Triton's autotune causes the Tensor._version to change, and so Pytorch autograd
+        # does a memcpy. To avoid this we run in inference_mode, which doesn't track the version.
+        with torch.inference_mode():
+            dq = torch.empty_like(q)
+            dk = torch.empty_like(k)
+            dv = torch.empty_like(v)
+            _flash_attn_backward(do, q, k, v, o, lse, dq, dk, dv,
+                                 bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale)
+        return dq, dk, dv, None, None, None
+
+
+flash_attn_func = FlashAttnFunc.apply

modeling.py

@@ -0,0 +1,255 @@
+import torch
+from torch import nn
+from dataclasses import dataclass
+from enum import Enum
+from typing import *
+from flash_attn import flash_attn_func
+from flash_attn_triton import flash_attn_func as flash_attn_func_triton
+from math import ceil
+
+
+class AttentionBackend(Enum):
+    Naive = 0
+    FlashAttentionCuda = 1
+    FlashAttentionTriton = 2
+
+
+global_config = {
+    'attn_backend': AttentionBackend.Naive
+}
+
+
+@dataclass
+class TransformerConfig:
+    vocab_size: int = -1,
+    num_layers: int = -1,
+    num_heads: int = -1,
+    hidden_size: int = -1,
+    max_seq_len: int = -1,
+    root_model: 'ToyTransformer' = None
+    device: torch.device = torch.device('cpu')
+    dtype: torch.dtype = torch.float32
+
+
+def expand_attn_mask(custom_attn_mask: torch.Tensor):
+    B, T = custom_attn_mask.shape
+    mask = custom_attn_mask.unsqueeze(1).repeat((1, T, 1))
+    seq_index_mask = (mask == custom_attn_mask[:, torch.arange(T)].view(B, T, 1))
+    return seq_index_mask & (torch.tril(mask) > 0)
+
+
+# expand attn mask to cu_seqlens for flash attn
+def expand_attn_mask_to_seq_lengths(attn_mask: torch.Tensor):
+    attn_mask = attn_mask.to('cpu')
+    seq_len = attn_mask.shape[0] * attn_mask.shape[1]
+    disjoint_point = torch.cat([torch.tensor([[True]] * attn_mask.shape[0]), attn_mask[:, 1:] != attn_mask[:, :-1]], dim=1)
+    return torch.cat([torch.nonzero(disjoint_point.view((-1,))), torch.tensor([[seq_len]])]).to(dtype=torch.int32)
+
+
+# naive RoPE implementation following https://arxiv.org/pdf/2104.09864.pdf
+def get_rope_cache_slow(seq_len: int, dim: int, theta: int, device: torch.device, dtype: torch.dtype):
+    assert dim % 2 == 0
+    freqs = theta ** (-2 * torch.arange(0, dim // 2, 1.) / dim)
+    freqs = torch.repeat_interleave(freqs, 2)
+    v1 = torch.cos(torch.arange(seq_len, dtype=torch.float).view((seq_len, 1)) * freqs)
+    v2 = torch.sin(torch.arange(seq_len, dtype=torch.float).view((seq_len, 1)) * freqs)
+    v2 = v2 * torch.tensor([1, -1] * (dim // 2))
+    indices = torch.tensor([j for i in range(0, dim, 2) for j in (i + 1, i)])
+    return v1.to(device, dtype=dtype), v2.to(device, dtype=dtype), indices.to(device)
+
+
+def apply_rope_slow(x, rope_cache, positions: Optional[torch.Tensor] = None):
+    v1, v2, indices = rope_cache
+    seq_len, dim = x.shape[1:]
+    if positions is None:
+        v1 = v1[:seq_len, :]
+        v2 = v2[:seq_len, :]
+    else:
+        v1 = v1[positions, torch.arange(dim)].view((-1, dim))
+        v2 = v2[positions, torch.arange(dim)].view((-1, dim))
+    applied_x = x * v1 + (x * v2)[:, :, indices]
+    return applied_x
+
+
+# Optimized RoPE implementation adapted from https://github.com/facebookresearch/llama/blob/main/llama/model.py
+def get_rope_cache_fast(seq_len: int, dim: int, theta: int, device: torch.device, dtype: torch.dtype):
+    freqs = (1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)))
+    t = torch.arange(seq_len, device=freqs.device)
+    freqs = torch.outer(t, freqs).float()
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+    return freqs_cis.to(device)
+
+
+def apply_rope_fast(x, rope_cache, positions: Optional[torch.Tensor] = None) -> torch.Tensor:
+    x_ = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+    if positions is None and x.shape[1] < rope_cache.shape[0]:
+        freqs_cis = rope_cache[:x.shape[1], :]
+    elif positions is not None:
+        freqs_cis = rope_cache[positions, :]
+    else:
+        freqs_cis = rope_cache
+    freqs_cis = freqs_cis.view([d if i == 1 or i == x_.ndim - 1 else 1 for i, d in enumerate(x_.shape)])
+
+    applied_x = torch.view_as_real(x_ * freqs_cis).flatten(2)
+    return applied_x.type_as(x)
+
+
+# RMSNorm implementation following https://arxiv.org/pdf/1910.07467.pdf
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, dtype, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor):
+        x_ = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return self.weight * x_
+
+
+class AttentionHead(nn.Module):
+    def __init__(self, config: TransformerConfig):
+        super().__init__()
+        self.config = config
+        self.head_size = config.hidden_size // config.num_heads
+        self.dtype = config.dtype
+        self.q_proj = nn.Linear(config.hidden_size, self.head_size, dtype=config.dtype)
+        self.k_proj = nn.Linear(config.hidden_size, self.head_size, dtype=config.dtype)
+        self.v_proj = nn.Linear(config.hidden_size, self.head_size, dtype=config.dtype)
+
+    def forward(self, x: torch.Tensor, attn_masked_bias: Optional[torch.Tensor],
+                kv_cache: Optional[List[torch.Tensor]]) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        q = self.q_proj(x)
+        k = self.k_proj(x)
+        v = self.v_proj(x)
+
+        # if global_config['attn_backend'] == AttentionBackend.FlashAttentionTriton:
+        # padding the position indices for alignment
+        # positions = torch.tensor([kv_cache[0].shape[1]] * q.shape[1]).to(q.device) if kv_cache is not None else torch.arange(0, x.shape[1], 1).to(q.device)
+
+        positions = torch.tensor([kv_cache[0].shape[1]]).to(q.device) if kv_cache is not None else None
+        q = apply_rope_fast(q, self.config.root_model.rope_cache, positions)
+        k = apply_rope_fast(k, self.config.root_model.rope_cache, positions)
+
+        if kv_cache is not None:
+            k = torch.concat([kv_cache[0], k], dim=1)
+            v = torch.concat([kv_cache[1], v], dim=1)
+
+        if global_config['attn_backend'] == AttentionBackend.FlashAttentionCuda:
+            q, k, v, = q.unsqueeze(2), k.unsqueeze(2), v.unsqueeze(2)
+            attn_result = flash_attn_func(q, k, v, causal=True)
+            q, k, v, attn_result = q.squeeze(2), k.squeeze(2), v.squeeze(2), attn_result.squeeze(2)
+        elif global_config['attn_backend'] == AttentionBackend.FlashAttentionTriton:
+            q, k, v, = q.unsqueeze(2), k.unsqueeze(2), v.unsqueeze(2)
+            attn_result = flash_attn_func_triton(q, k, v, attn_masked_bias.unsqueeze(1) if attn_masked_bias is not None else None,
+                                                 True if kv_cache is None else False)
+            q, k, v, attn_result = q.squeeze(2), k.squeeze(2), v.squeeze(2), attn_result.squeeze(2)
+        else:
+            attn_score = (q @ k.permute(0, 2, 1) / (self.head_size ** 0.5)) + attn_masked_bias
+            attn_result = torch.softmax(attn_score, dim=2) @ v
+
+        return attn_result, [k, v]
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, config: TransformerConfig):
+        super().__init__()
+        self.config = config
+        self.attn_heads = nn.ModuleList([AttentionHead(config) for _ in range(config.num_heads)])
+        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, dtype=config.dtype)
+
+    def forward(self, x: torch.Tensor, attn_masked_bias: Optional[torch.Tensor],
+                kv_cache: Optional[List[torch.Tensor]]) -> Tuple[torch.Tensor, List[List[torch.Tensor]]]:
+        head_outputs = [head(x, attn_masked_bias, kv_cache[idx] if kv_cache is not None else None) for idx, head in
+                        enumerate(self.attn_heads)]
+        return self.o_proj(torch.concat([o[0] for o in head_outputs], dim=2)), [o[1] for o in head_outputs]
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, config: TransformerConfig):
+        super().__init__()
+        self.config = config
+        self.mha = MultiHeadAttention(config)
+        self.up_proj = nn.Linear(config.hidden_size, config.hidden_size * 4, dtype=config.dtype)
+        self.down_proj = nn.Linear(config.hidden_size * 4, config.hidden_size, dtype=config.dtype)
+        self.ln_mha = nn.LayerNorm(config.hidden_size, dtype=config.dtype)
+        self.ln_ffn = nn.LayerNorm(config.hidden_size, dtype=config.dtype)
+        self.act = nn.GELU()
+
+    def forward(self, x: torch.Tensor, attn_masked_bias: Optional[torch.Tensor],
+                kv_cache: Optional[List[torch.Tensor]]) -> Tuple[torch.Tensor, List[List[torch.Tensor]]]:
+        mha_output, new_kv_cache = self.mha(self.ln_mha(x), attn_masked_bias, kv_cache)
+        mha_output = x + mha_output
+        ffn_output = self.down_proj(self.act(self.up_proj(self.ln_ffn(mha_output))))
+        return mha_output + ffn_output, new_kv_cache
+
+
+class ToyTransformer(nn.Module):
+    def __init__(self, vocab_size: int, num_layers: int, num_heads: int, hidden_size: int, max_seq_len: int,
+                 device: torch.device = torch.device('cpu'), dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.config = TransformerConfig(vocab_size, num_layers, num_heads, hidden_size, max_seq_len, self, device,
+                                        dtype)
+
+        self.sem_embed = nn.Embedding(vocab_size, hidden_size, dtype=dtype)
+
+        self.rope_cache = get_rope_cache_fast(max_seq_len, hidden_size // num_heads, 10000, device, dtype)
+
+        self.decoder_layers = nn.ModuleList([DecoderLayer(self.config) for _ in range(num_layers)])
+        self.lm_head = nn.Linear(hidden_size, vocab_size, dtype=dtype)
+        self.to(device)
+
+    def forward(self, seq: torch.Tensor,
+                attn_mask: Optional[torch.Tensor] = None,
+                kv_cache: Optional[List[torch.Tensor]] = None) -> Tuple[torch.Tensor, List[List[List[torch.Tensor]]]]:
+        # sanity checks
+        assert attn_mask is None or kv_cache is None  # No support for attn_mask and kv_cache both enabled
+        if kv_cache is not None:
+            assert seq.shape[0] == 1, 'kv_cache is not supported for batch inference'
+        # handle flash-attn triton alignment requirement (actually only needed for backward)
+        seq_length = seq.shape[1]
+        if kv_cache is None and global_config['attn_backend'] == AttentionBackend.FlashAttentionTriton and seq_length % 128 != 0:
+            if attn_mask is None:  # forcibly enable attn_mask due to padding
+                attn_mask = torch.ones(seq.shape, device=self.device)
+            pad_length = (ceil(seq_length / 128) * 128) - seq_length
+            seq = nn.functional.pad(seq, (0, pad_length))
+            attn_mask = nn.functional.pad(attn_mask, (0, pad_length))
+
+        # handle attn_bias
+        if global_config['attn_backend'] == AttentionBackend.FlashAttentionCuda:
+            assert attn_mask is None, 'FlashAttn-Cuda does not support custom attn_mask'
+            attn_masked_bias = None
+        elif global_config['attn_backend'] == AttentionBackend.FlashAttentionTriton and attn_mask is None:
+            attn_masked_bias = None
+        elif attn_mask is not None:
+            attn_masked_bias = expand_attn_mask(attn_mask)
+        elif attn_mask is None and kv_cache is None:
+            attn_masked_bias = expand_attn_mask(torch.ones(seq.shape, device=self.device))
+        elif kv_cache is not None:
+            attn_masked_bias = torch.ones((1, seq.shape[1], seq.shape[1]), dtype=torch.bool, device=self.device)
+        else:
+            attn_masked_bias = None
+
+        if attn_masked_bias is not None:
+            mask_zero = torch.tensor(0, dtype=self.config.dtype)
+            mask_val = torch.tensor(torch.finfo(self.config.dtype).min / 2, dtype=self.config.dtype)
+            attn_masked_bias = torch.where(attn_masked_bias, mask_zero, mask_val).to(self.device)
+
+        hidden = self.sem_embed(seq)
+
+        new_kv_cache = []
+        for idx, decoder in enumerate(self.decoder_layers):
+            hidden, layer_kv_cache = decoder(hidden, attn_masked_bias, kv_cache[idx] if kv_cache is not None else None)
+            new_kv_cache.append(layer_kv_cache)
+
+        logits = self.lm_head(hidden)
+
+        # remove padding for flash-attn triton
+        if kv_cache is None and global_config['attn_backend'] == AttentionBackend.FlashAttentionTriton and seq_length % 128 != 0:
+            logits = logits[:, :seq_length, :]
+            new_kv_cache = [[[cache[:, :seq_length, :] for cache in head] for head in layer] for layer in new_kv_cache]
+
+        return logits, new_kv_cache
+
+    @property
+    def device(self):
+        return next(self.parameters()).device

tokenizers.py

@@ -0,0 +1,244 @@
+import time
+from typing import *
+import re
+import json
+import numba
+
+
+def sample_vocab(tokens: Iterable[str], vocab_size: Optional[int] = None,
+                 vocab_coverage: Optional[float] = None) -> List[str]:
+    assert (vocab_size is not None and vocab_coverage is None) or \
+           (vocab_size is None and vocab_coverage is not None), "vocab_size [or] vocab_coverage need specified"
+
+    token_count = {}
+    for c in tokens:
+        token_count[c] = token_count.get(c, 0) + 1
+
+    if vocab_size is not None:
+        token_count = list(token_count.items())
+        token_count.sort(key=lambda i: i[1], reverse=True)
+        vocab = [c[0] for c in token_count[:vocab_size]]
+    else:
+        total_count = sum(token_count.values())
+        token_freq = [(c, i / total_count) for c, i in token_count.items()]
+        token_freq.sort(key=lambda i: i[1], reverse=True)
+        freq_sum = 0.0
+        split = 0
+        for split in range(len(token_freq)):
+            freq_sum += token_freq[split][1]
+            if freq_sum >= vocab_coverage:
+                break
+        vocab = [c[0] for c in token_freq[:split + 1]]
+    return vocab
+
+
+class CharTokenizer:
+    def __init__(self, corpus: str, vocab_size: Optional[int] = None, vocab_coverage: Optional[float] = None,
+                 reserved_vocab: Optional[List[str]] = None, unk_literal: str = '<unk>'):
+        if reserved_vocab is not None:
+            assert len(reserved_vocab) == len(set(reserved_vocab)), 'no duplicate is allowed in reserved vocab'
+            assert unk_literal not in reserved_vocab, f'unk literal "{unk_literal}" cannot be in reserved vocab'
+        else:
+            reserved_vocab = []
+        vocab = reserved_vocab.copy() if reserved_vocab is not None else []
+        vocab += sample_vocab(corpus, vocab_size - len(vocab) - 1, vocab_coverage)
+        self.s2i = {s: i + 1 for i, s in enumerate(vocab)}
+        self.s2i[unk_literal] = 0
+        self.i2s = {i: s for s, i in self.s2i.items()}
+        self.special_vocab = set(reserved_vocab + [unk_literal])
+        self.unk_literal = unk_literal
+
+    def encode(self, text: str) -> List[int]:
+        cursor, ids = 0, []
+        while cursor < len(text):
+            for s in self.special_vocab:
+                if text[cursor:].startswith(s):
+                    ids.append(self.s2i[s])
+                    cursor += len(s)
+                    break
+            else:
+                ids.append(self.s2i.get(text[cursor], self.s2i.get(self.unk_literal)))
+                cursor += 1
+        return ids
+
+    def decode(self, ids: List[int]) -> str:
+        return ''.join(self.i2s[i] for i in ids)
+
+    def get_vocab_mapping(self):
+        return self.s2i
+
+
+class WordTokenizer:
+    def __init__(self, corpus: str, vocab_size: Optional[int] = None, vocab_coverage: Optional[float] = None,
+                 reserved_vocab: Optional[List[str]] = None, unk_literal: str = '<unk>'):
+        if reserved_vocab is not None:
+            assert len(reserved_vocab) == len(set(reserved_vocab)), 'no duplicate is allowed in reserved vocab'
+            assert unk_literal not in reserved_vocab, f'unk literal "{unk_literal}" cannot be in reserved vocab'
+        else:
+            reserved_vocab = []
+        vocab = reserved_vocab.copy() if reserved_vocab is not None else []
+
+        tokens = (c[0] if c[0] != '' else c[1] for c in re.finditer(r'(\w+)|(\W)', corpus))
+        vocab += sample_vocab(tokens, vocab_size - len(vocab) - 1, vocab_coverage)
+
+        self.s2i = {s: i + 1 for i, s in enumerate(vocab)}
+        self.s2i[unk_literal] = 0
+        self.i2s = {i: s for s, i in self.s2i.items()}
+        self.special_vocab = set(reserved_vocab + [unk_literal])
+        self.unk_literal = unk_literal
+
+    def encode(self, text: str) -> List[int]:
+        specials = '|'.join(f'{i}' for i in self.special_vocab)
+        tokens = (c[0] if c[0] != '' else c[1] for c in re.finditer(rf'({specials}|\w+)|(\W)', text))
+        return [self.s2i.get(t, self.s2i[self.unk_literal]) for t in tokens]
+
+    def decode(self, ids: List[int]) -> str:
+        return ''.join(self.i2s[i] for i in ids)
+
+    def get_vocab_mapping(self):
+        return self.s2i
+
+    def get_vocab_size(self):
+        return len(self.s2i)
+
+    def eval_vocab_coverage(self, corpus: str):
+        encoded = self.encode(corpus)
+        return 1 - (len([i for i in encoded if i == 0]) / len(encoded))
+
+
+class TRIETokenizer:
+    @staticmethod
+    def split_bytes(data: bytes):
+        return [b'%c' % i for i in data]
+
+    def __init__(self, vocab_file: str):
+        self.nodes = [(b'', -1, -1, [-1 for _ in range(256)])]  # node value, parent index, token id, children
+        with open(vocab_file, 'r') as file:
+            vocabs = json.load(file)
+        vocabs.sort(key=lambda i: len(i['bytes']))
+        for entry in vocabs:
+            self.add_vocab(bytes(entry['bytes']), entry['id'])
+
+        self.id_to_bytes = {i['id']: i['bytes'] for i in vocabs}
+
+    def add_vocab(self, vocab_bytes: bytes, vocab_id: int):
+        cur_node_idx = 0
+        for i, b in enumerate(vocab_bytes):
+            cur_node = self.nodes[cur_node_idx]
+            if cur_node[3][b] != -1:
+                cur_node_idx = cur_node[3][b]
+            else:
+                new_node_idx = len(self.nodes)
+                self.nodes.append((vocab_bytes, cur_node_idx, vocab_id if i == len(vocab_bytes) - 1 else -1,
+                                   [-1 for _ in range(256)]))
+                cur_node[3][b] = new_node_idx
+                cur_node_idx = new_node_idx
+
+    def attempt_match(self, match_bytes: bytes):
+        match_length, match_token_id = -1, -1
+        cur_node_idx, depth = 0, 0
+        for i, b in enumerate(match_bytes):
+            match_node_idx = self.nodes[cur_node_idx][3][b]
+            if match_node_idx == -1:
+                break
+            cur_node = self.nodes[match_node_idx]
+            if cur_node[2] != -1:
+                match_length = depth
+                match_token_id = cur_node[2]
+            cur_node_idx = match_node_idx
+            depth += 1
+        return match_length, match_token_id
+
+    def encode(self, text: str):
+        text_bytes = text.encode('utf-8')
+        tokens, length = [], 0
+        while length < len(text_bytes):
+            offset, token_id = self.attempt_match(text_bytes[length:])
+            assert offset >= 0
+            tokens.append(token_id)
+            length += offset + 1
+        return tokens
+
+    def decode(self, token_ids: List[int]):
+        return bytes([t for i in token_ids for t in self.id_to_bytes[i]]).decode('utf-8', errors='replace')
+
+    def get_vocab_size(self):
+        return len(self.id_to_bytes)
+
+
+@numba.njit
+def trie_attempt_match_jit(trie_nodes, match_bytes: bytes):
+    match_length, match_token_id = -1, -1
+    cur_node_idx, depth = 0, 0
+    for i, b in enumerate(match_bytes):
+        match_node_idx = trie_nodes[cur_node_idx][3][int(b)]
+        if match_node_idx == -1:
+            break
+        cur_node = trie_nodes[match_node_idx]
+        if cur_node[2] != -1:
+            match_length = depth
+            match_token_id = cur_node[2]
+        cur_node_idx = match_node_idx
+        depth += 1
+    return match_length, match_token_id
+
+
+@numba.njit
+def trie_encode_jit(trie_nodes, text_bytes: bytes):
+    tokens, length = [], 0
+    while length < len(text_bytes):
+        offset, token_id = trie_attempt_match_jit(trie_nodes, text_bytes[length:])
+        assert offset >= 0
+        tokens.append(token_id)
+        length += offset + 1
+    return tokens
+
+
+class TRIETokenizerFast:
+    def __init__(self, vocab_file: str):
+        self.nodes = [(b'', -1, -1, [-1 for _ in range(256)])]  # node value, parent index, token id, children
+        with open(vocab_file, 'r') as file:
+            vocabs = json.load(file)
+        vocabs.sort(key=lambda i: len(i['bytes']))
+        for entry in vocabs:
+            self.add_vocab(bytes(entry['bytes']), entry['id'])
+
+        self.id_to_bytes = {i['id']: i['bytes'] for i in vocabs}
+
+        self.nodesJit = numba.typed.List(self.nodes)
+
+    def add_vocab(self, vocab_bytes: bytes, vocab_id: int):
+        cur_node_idx = 0
+        for i, b in enumerate(vocab_bytes):
+            cur_node = self.nodes[cur_node_idx]
+            if cur_node[3][b] != -1:
+                cur_node_idx = cur_node[3][b]
+            else:
+                new_node_idx = len(self.nodes)
+                self.nodes.append((vocab_bytes, cur_node_idx, vocab_id if i == len(vocab_bytes) - 1 else -1,
+                                   [-1 for _ in range(256)]))
+                cur_node[3][b] = new_node_idx
+                cur_node_idx = new_node_idx
+
+    def encode(self, text: str):
+        return trie_encode_jit(self.nodesJit, text.encode('utf-8'))
+
+    def decode(self, token_ids: List[int]):
+        return bytes([t for i in token_ids for t in self.id_to_bytes[i]]).decode('utf-8', errors='replace')
+
+    def get_vocab_size(self):
+        return len(self.id_to_bytes)
+
+# if __name__ == '__main__':
+#     tokenizer = TRIETokenizerFast('llama_vocab_pruned_20k.json')
+#     with open('corpus/TinyStoriesV2-GPT4-valid.txt', 'r') as file:
+#         text = file.read()[:10240]
+#
+#     total_tokens = 0
+#     s = time.time()
+#     for i in range(1000):
+#         encoded = tokenizer.encode(text)
+#         total_tokens += len(encoded)
+#         print(len(encoded))
+#     e = time.time()
+#     print(f'{e - s:.3f} secs, {total_tokens / (e - s):.3f} tps')

webui.py

@@ -0,0 +1,142 @@
+import gradio as gr
+from modeling import global_config, ToyTransformer, AttentionBackend
+import torch
+from tokenizers import TRIETokenizer
+from threading import Thread
+import bisect
+
+if torch.cuda.is_available():
+    g_device = torch.device('cpu')
+else:
+    g_device = torch.device('cpu')
+global_config['attn_backend'] = AttentionBackend.Naive
+
+g_SEQ_LEN = 1024
+g_HIDDEN_SIZE = 768
+g_NUM_HEADS = 12
+g_NUM_LAYERS = 12
+g_DTYPE = torch.float32
+
+g_tokenizer = TRIETokenizer('llama_vocab_pruned_32k.json')
+g_model = ToyTransformer(g_tokenizer.get_vocab_size(), g_NUM_LAYERS, g_NUM_HEADS, g_HIDDEN_SIZE, g_SEQ_LEN, g_device, g_DTYPE)
+
+g_model.load_state_dict(torch.load('model.pt', map_location='cpu'))
+
+
+def generate(model, tokenizer, prompt, temperature, top_p, rep_penalty,
+             max_new_tokens=20, total_tokens=None,
+             end_tokens=None,
+             enable_kv_cache=True):
+    model.eval()
+
+    feed_tokens = tokenizer.encode(prompt) if isinstance(prompt, str) else prompt
+
+    all_tokens = feed_tokens.copy()
+    if total_tokens is not None:
+        max_new_tokens = max(0, total_tokens - len(feed_tokens))
+
+    with torch.no_grad():
+        kv_cache = None
+        for _ in range(max_new_tokens):
+            logits, kv_cache = model.forward(
+                torch.tensor([feed_tokens if enable_kv_cache else all_tokens]).to(model.device),
+                kv_cache=kv_cache)
+            logits = logits[0][-1].cpu()
+            if not enable_kv_cache:
+                kv_cache = None
+
+            # apply repetition penalty
+            logits_rep = torch.gather(logits, 0, torch.tensor(all_tokens))
+            logits_rep = torch.where(logits_rep < 0, logits_rep * rep_penalty, logits_rep / rep_penalty)
+            logits.scatter_(0, torch.tensor(all_tokens), logits_rep)
+
+            # apply temperature
+            logits /= max(temperature, 1e-6)
+
+            probs = torch.softmax(logits, dim=0)
+
+            # apply top-p
+            ordered_probs, ordered_indices = torch.sort(probs, descending=True)
+            cum_probs = torch.cumsum(ordered_probs, dim=0).tolist()
+            top_p_index = bisect.bisect_right(cum_probs, top_p) + 1
+            ordered_probs, ordered_indices = ordered_probs[:top_p_index], ordered_indices[:top_p_index]
+            sampled_index = ordered_indices[torch.multinomial(ordered_probs, num_samples=1).item()].item()
+
+            all_tokens.append(sampled_index)
+            feed_tokens = [sampled_index]
+
+            if end_tokens is not None and sampled_index in end_tokens:
+                break
+
+            yield feed_tokens
+    return
+
+
+def predict(user_input, history, max_length, top_p, temperature, rep_penalty, retry):
+    if retry and len(history) == 0:
+        yield []
+        return
+    elif retry:
+        user_input = history[-1][0]
+        history = history[:-1]
+
+    history.append((user_input, ""))
+
+    encoded_inputs = [(g_tokenizer.encode('User:' + h[0]), g_tokenizer.encode('Assistant:' + h[1])) for h in history]
+    taken_rounds, taken_rounds_length = [], 0
+    while len(taken_rounds) < len(encoded_inputs):
+        round_pair = encoded_inputs[len(encoded_inputs) - 1 - len(taken_rounds)]
+        if len(round_pair[0]) + len(round_pair[1]) + taken_rounds_length >= g_SEQ_LEN - max_length:
+            break
+        taken_rounds.append(round_pair)
+        taken_rounds_length += len(round_pair[0]) + len(round_pair[1])
+    taken_rounds = taken_rounds[::-1]
+
+    input_tokens = g_tokenizer.encode('<s>A chat between User and Assistant.')
+    for round_pair in taken_rounds:
+        input_tokens += g_tokenizer.encode('\n') + round_pair[0] + g_tokenizer.encode('\n') + round_pair[1]
+    # print(taken_rounds, g_tokenizer.decode(input_tokens))
+    for response in generate(g_model, g_tokenizer, input_tokens, temperature, top_p, rep_penalty, max_length, end_tokens=g_tokenizer.encode('</s>')):
+        history[-1] = (history[-1][0], history[-1][1] + g_tokenizer.decode(response))
+        yield history
+
+
+def main():
+    css = '''
+        .contain {max-width:50}
+
+        #chatbot {min-height:500px}
+    '''
+
+    with gr.Blocks(css=css) as demo:
+        gr.HTML('<h1 align="center">ToyTransformer</h1>')
+
+        chatbot = gr.Chatbot(elem_id='chatbot')
+        with gr.Column():
+            user_input = gr.Textbox(show_label=False, placeholder="输入", lines=1, container=False)
+            with gr.Row():
+                submitBtn = gr.Button("Send", variant="primary")
+                retryBtn = gr.Button("Retry")
+                cancelBtn = gr.Button('Undo')
+                emptyBtn = gr.Button("Clear")
+            with gr.Row():
+                max_length = gr.Slider(0, 512, value=200, step=1, label="Max Response Tokens", interactive=True)
+                top_p = gr.Slider(0, 1, value=0.7, step=0.01, label="Top-P", interactive=True)
+                temperature = gr.Slider(0, 1, value=0.5, step=0.01, label="Temperature", interactive=True)
+                rep_penalty = gr.Slider(1.0, 1.5, value=1.1, step=0.01, label='Repetition Penalty', interactive=True)
+
+        submitBtn.click(predict, [user_input, chatbot, max_length, top_p, temperature, rep_penalty, gr.State(False)],
+                        [chatbot], show_progress=False)
+        submitBtn.click(lambda: '', [], [user_input], show_progress=False)
+
+        retryBtn.click(predict, [user_input, chatbot, max_length, top_p, temperature, rep_penalty, gr.State(True)],
+                       [chatbot], show_progress=False)
+
+        cancelBtn.click(lambda m: m[:-1], [chatbot], [chatbot], show_progress=False)
+
+        emptyBtn.click(lambda: [], outputs=[chatbot], show_progress=False)
+
+    demo.queue().launch(share=False, inbrowser=True)
+
+
+main()