honglu2875

"""Contains generation/sampling code"""

# This file contains code from https://github.com/facebookresearch/llama which
# is available under the following licence:

# Copyright (c) Meta Platforms, Inc. and affiliates.
# This software may be used and distributed according to the terms of the GNU
# General Public License version 3.

import torch

honglu2875

_

honglu2875

import torch
from torch import nn
import time

bias = torch.tril(torch.ones((2048, 2048), dtype=torch.uint8, device='cuda')).view(
                1, 1, 2048, 2048
            )
def _attn(
    query,
    key,