Birch-san

from __future__ import annotations
from typing import Any, Literal, NamedTuple, Optional
from os import environ
import triton
import triton.language as tl
import torch
from torch import Tensor, enable_grad
from torch.autograd import Function
from torch.autograd.function import FunctionCtx
import torch.autograd.forward_ad as fwAD

Birch-san

from __future__ import annotations
from typing import NamedTuple, Optional
from argparse import ArgumentParser, Namespace
from dataclasses import dataclass
import torch
from torch import Tensor, inference_mode
from torch.nn import Module, Linear
from torch.nn.functional import relu
from einops import rearrange

Birch-san

from __future__ import annotations
"""
Fused Attention
===============

This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao (https://tridao.me/publications/flash2/flash2.pdf)

Credits: OpenAI kernel team
Extra Credits:

Birch-san

from __future__ import annotations
from argparse import ArgumentParser, Namespace
from dataclasses import dataclass
from functools import partial
from os import environ
from typing import Any, Callable, NamedTuple, Optional
import torch
from torch import Tensor, no_grad, enable_grad
from torch.autograd import Function
from torch.autograd.function import FunctionCtx

Birch-san

"""
Fused Attention
===============

This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao (https://tridao.me/publications/flash2/flash2.pdf)

Credits: OpenAI kernel team

Extra Credits:

Birch-san

from dataclasses import dataclass
from einops import rearrange
import re
import torch
from torch import BoolTensor, FloatTensor, IntTensor, LongTensor, inference_mode
from torch.nn.functional import pad
from itertools import islice
from typing import Generator, Iterable, Iterator, Optional, Protocol, TypeVar
from typing_extensions import override
from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel, UNet2DConditionOutput

Birch-san

from __future__ import annotations
from argparse import ArgumentParser, Namespace
from dataclasses import dataclass
from functools import partial
from typing import Any, Callable, Optional
import torch
from torch import Tensor, no_grad, enable_grad
import torch.autograd.forward_ad as fwAD
from torch.nn.attention import SDPBackend, sdpa_kernel
from torch.nn.functional import scaled_dot_product_attention

Birch-san

from __future__ import annotations
from argparse import ArgumentParser, Namespace
from dataclasses import dataclass
from functools import partial
from typing import Callable, Generic, TypeVar
import torch
from torch import enable_grad, no_grad
import torch.autograd.forward_ad as fwAD
from torch.func import linearize
from torch.nn.attention import SDPBackend, sdpa_kernel

Birch-san

from __future__ import annotations
from argparse import ArgumentParser, Namespace
from dataclasses import dataclass
from functools import partial
from typing import Callable
import torch
from torch import enable_grad, no_grad
import torch.autograd.forward_ad as fwAD
from torch.nn.attention import SDPBackend, sdpa_kernel
from torch.nn.functional import scaled_dot_product_attention

Birch-san

#!/usr/bin/env bash
set -eo pipefail

# https://stackoverflow.com/a/12194427/5257399
create() { # fd base [qualifier [suffix [max]]]
  local fd="$1" base="$2" qualifier="${3-}" suffix="${4-.png}" max="${5-}"
  local n=0 file
  local - # ash-style local scoping of options in 4.4+
  set -o noclobber
  REPLY=