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 -*-

import torch
import triton
import triton.language as tl


@triton.jit
def cumsum_matmul_kernel(
    s,

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 -*-
# 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

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) 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

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

import glob
import orjson
import os

import datasets
from itertools import islice

_DESCRIPTION = """