yzhangcs

# -*- coding: utf-8 -*-

import glob
import orjson
import os

import datasets
from itertools import islice

_DESCRIPTION = """

yzhangcs

# -*- coding: utf-8 -*-

# Copyright (c) 2023, Tri Dao.
# https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/ops/triton/rotary.py

from typing import Optional, Tuple, Union

import torch
import torch.nn as nn
import triton

yzhangcs

sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)"
git clone https://github.com/Pilaton/OhMyZsh-full-autoupdate.git ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/ohmyzsh-full-autoupdate
git clone https://github.com/zsh-users/zsh-autosuggestions ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-autosuggestions
git clone https://github.com/zsh-users/zsh-completions ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-completions
git clone https://github.com/zsh-users/zsh-history-substring-search ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-history-substring-search
git clone https://github.com/zsh-users/zsh-syntax-highlighting ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-syntax-highlighting

yzhangcs

# -*- coding: utf-8 -*-
# Copyright (c) 2024, Songlin Yang, Yu Zhang

from typing import Optional, Tuple

import torch
import triton
import triton.language as tl

from fla.ops.common.chunk_h_split import chunk_bwd_dh, chunk_fwd_h

yzhangcs

# -*- coding: utf-8 -*-
# Copyright (c) 2024, Songlin Yang, Yu Zhang

from typing import Optional, Tuple

import torch
import triton
import triton.language as tl

yzhangcs

# -*- coding: utf-8 -*-

import torch
import triton
import triton.language as tl


@triton.jit
def cumsum_matmul_kernel(
    s,

yzhangcs

# -*- coding: utf-8 -*-

import torch
import torch.nn as nn
import triton
import triton.language as tl
from transformers.models.llama.modeling_llama import LlamaRMSNorm


@triton.jit

yzhangcs

# -*- coding: utf-8 -*-

from typing import Optional, Tuple

import torch
import torch.nn as nn
import triton
import triton.language as tl
from transformers.models.llama.modeling_llama import (LlamaAttention,
                                                      apply_rotary_pos_emb,

yzhangcs

# -*- coding: utf-8 -*-

from __future__ import annotations

import copy
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

yzhangcs

# Terminal color definitions

class fg:
    BLACK   = '\033[30m'
    RED     = '\033[31m'
    GREEN   = '\033[32m'
    YELLOW  = '\033[33m'
    BLUE    = '\033[34m'
    MAGENTA = '\033[35m'
    CYAN    = '\033[36m'