examples/flex_attn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FlexAttention API Usage Notebook\n",
    "\n",
    "This notebook demonstrates the usage of the new FlexAttention API, which allows users to specify modifications to the computed attention scores in Scaled Dot Product Attention (SDPA).\n",
    "\n",
    "## Table of Contents\n",
    "1. [Introduction](#introduction)\n",
    "2. [Setup](#setup)\n",
    "3. [Basic Usage](#basic-usage)\n",
    "4. [Score Modification vs Score Masking](#Score-Modification-vs-Score-Masking)\n",
    "4. [Score Modification Examples](#score-modification-examples)\n",
    "   - [Full Attention (No-op)](#full-attention-no-op)\n",
    "   - [Standard Causal Mask](#standard-causal-mask)\n",
    "   - [Sliding Window Attention](#sliding-window-attention)\n",
    "   - [Prefix LM (Bidirectional + Causal)](#prefix-lm-bidirectional--causal)\n",
    "   - [Document Masking](#document-masking)\n",
    "   - [NATTEN Masking](#natten-masking)\n",
    "   - [Alibi Bias](#alibi-bias)\n",
    "   - [Tanh Soft-Capping](#tanh-soft-capping)\n",
    "   - [Nested Jagged Tensor](#nested-jagged-tensor)\n",
    "   - [Flamingo Cross Attention](#flamingo-cross-attention)\n",
    "\n",
    "## Introduction\n",
    "\n",
    "The FlexAttention API allows users to specify custom modifications to attention scores within a Fused Scaled Dot Product Attention Kernel. This enables various attention patterns and biases to be implemented efficiently, with potential runtime and memory savings. The API will also generate fused backward kernels based off of the user defined modification.\n",
    "\n",
    "## Setup\n",
    "\n",
    "First, let's import the necessary libraries and set up our environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "output": {
     "id": 2064820487248880,
     "loadingStatus": "loaded"
    }
   },
   "outputs": [],
   "source": [
    "import random\n",
    "from functools import lru_cache, partial\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from tabulate import tabulate\n",
    "from torch.nn.attention.flex_attention import (\n",
    "    _DEFAULT_SPARSE_BLOCK_SIZE,\n",
    "    create_block_mask,\n",
    "    create_mask,\n",
    "    flex_attention,\n",
    ")\n",
    "from triton.testing import do_bench\n",
    "\n",
    "torch.set_default_device(\"cuda\")\n",
    "torch.manual_seed(0)\n",
    "\n",
    "torch._dynamo.config.cache_size_limit = 1000\n",
    "\n",
    "# Compile the flex_attention function\n",
    "flex_attention = torch.compile(flex_attention, dynamic=False)\n",
    "\n",
    "# For better performance, you can use:\n",
    "# flex_attention = torch.compile(_flex_attention, dynamic=False, mode=\"max-autotune-no-cudagraphs\")\n",
    "\n",
    "data_type = torch.float16\n",
    "\n",
    "# The kernels will utilize block sparisty to increase performance\n",
    "print(f\"Using the default sparsity block size: {_DEFAULT_SPARSE_BLOCK_SIZE}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will define some helpful testing utilities that will print a block sparse representation of the score_mod function and mask_fn. \n",
    "\n",
    "As well, it will compare the performance between \n",
    "- FlexAttention \n",
    "- One of the SOTA implementation FlashAttentionV2 with Causal masking.\n",
    "- nn.F.scaled_dot_product_attention + fully materialized attn_mask. This will dispatch to a fused implementation `EFFICIENT_ATTENTION` that allows for arbitrary masking. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@lru_cache\n",
    "def create_block_mask_cached(score_mod, B, H, M, N, device=\"cuda\"):\n",
    "    block_mask = create_block_mask(score_mod, B, H, M, N, device=device)\n",
    "    return block_mask\n",
    "\n",
    "\n",
    "def calculate_tflops(flops: float, time_ms: float, multiplier: int) -> float:\n",
    "    return multiplier * flops * (1e3 / time_ms) / 1e12\n",
    "\n",
    "\n",
    "def test_mask(\n",
    "    score_mod=None,\n",
    "    mask_mod=None,\n",
    "    B=16,\n",
    "    H=16,\n",
    "    S=8192,\n",
    "    D=64,\n",
    "    skip_correctness=False,\n",
    "    print_mask=True,\n",
    "):\n",
    "    assert (\n",
    "        score_mod is not None or mask_mod is not None\n",
    "    ), \"Must provide a score_mod or mask_mod\"\n",
    "    query = torch.randn(\n",
    "        B, H, S, D, device=\"cuda\", dtype=torch.float16, requires_grad=True\n",
    "    )\n",
    "    key = torch.randn(\n",
    "        B, H, S, D, device=\"cuda\", dtype=torch.float16, requires_grad=True\n",
    "    )\n",
    "    value = torch.randn(\n",
    "        B, H, S, D, device=\"cuda\", dtype=torch.float16, requires_grad=True\n",
    "    )\n",
    "    gradOut = torch.randn(B, H, S, D, device=\"cuda\", dtype=torch.float16)\n",
    "\n",
    "    if mask_mod is not None:\n",
    "        block_mask = create_block_mask_cached(mask_mod, 1, 1, S, S, device=query.device)\n",
    "    else:\n",
    "        block_mask = None\n",
    "    sdpa_mask_fn = mask_mod if mask_mod is not None else score_mod\n",
    "    mask = create_mask(sdpa_mask_fn, 1, 1, S, S, device=query.device)\n",
    "\n",
    "    causal_fa2 = lambda: F.scaled_dot_product_attention(\n",
    "        query, key, value, is_causal=True\n",
    "    )\n",
    "    xformers_mask = lambda: F.scaled_dot_product_attention(\n",
    "        query, key, value, attn_mask=mask\n",
    "    )\n",
    "    flex_attention_call = lambda: flex_attention(\n",
    "        query, key, value, score_mod=score_mod, block_mask=block_mask\n",
    "    )\n",
    "\n",
    "    results = []\n",
    "    if block_mask is not None:\n",
    "        density = (100 - block_mask.sparsity()) / 100\n",
    "    else:\n",
    "        density = 1.0\n",
    "    causal_fav2_flops = 0.5 * B * H * D * S * S\n",
    "    flops = density * B * H * D * S * S\n",
    "\n",
    "    # Forward pass\n",
    "    causal_fa2_time = do_bench(causal_fa2)\n",
    "    xformers_mask_time = do_bench(xformers_mask)\n",
    "    flex_ms = do_bench(flex_attention_call)\n",
    "\n",
    "    # Backward pass\n",
    "    causal_fa2_out = causal_fa2()\n",
    "    xformers_out = xformers_mask()\n",
    "    flex_out = flex_attention_call()\n",
    "\n",
    "    causal_fa2_bw_time = do_bench(\n",
    "        lambda: causal_fa2_out.backward(gradOut, retain_graph=True)\n",
    "    )\n",
    "    xformers_mask_bw_time = do_bench(\n",
    "        lambda: xformers_out.backward(gradOut, retain_graph=True)\n",
    "    )\n",
    "    flex_bw_ms = do_bench(lambda: flex_out.backward(gradOut, retain_graph=True))\n",
    "\n",
    "    # Inline correctness check\n",
    "    if not skip_correctness:\n",
    "        xformers_outs = []\n",
    "        flex_outs = []\n",
    "\n",
    "        query.grad = None\n",
    "        key.grad = None\n",
    "        value.grad = None\n",
    "\n",
    "        out1 = xformers_mask()\n",
    "        xformers_outs.append(out1)\n",
    "        out1.backward(gradOut)\n",
    "        xformers_outs += [query.grad, key.grad, value.grad]\n",
    "\n",
    "        query.grad = None\n",
    "        key.grad = None\n",
    "        value.grad = None\n",
    "\n",
    "        out2 = flex_attention_call()\n",
    "        flex_outs.append(out2)\n",
    "        out2.backward(gradOut)\n",
    "        flex_outs += [query.grad, key.grad, value.grad]\n",
    "        for flex, xformer in zip(flex_outs, xformers_outs):\n",
    "            torch.testing.assert_close(flex, xformer, atol=1e-1, rtol=1e-2)\n",
    "\n",
    "        print(\"Correctness check passed ✅\")\n",
    "    # Usage in your results formatting:\n",
    "    results = [\n",
    "        [\n",
    "            \"causal FA2\",\n",
    "            f\"{causal_fa2_time:.4f}\",\n",
    "            f\"{calculate_tflops(causal_fav2_flops, causal_fa2_time, 4):.2f}\",\n",
    "            f\"{causal_fa2_bw_time:.4f}\",\n",
    "            f\"{calculate_tflops(causal_fav2_flops, causal_fa2_bw_time, 10):.2f}\",\n",
    "        ],\n",
    "        [\n",
    "            \"F.sdpa + mask\",\n",
    "            f\"{xformers_mask_time:.4f}\",\n",
    "            f\"{calculate_tflops(flops, xformers_mask_time, 4):.2f}\",\n",
    "            f\"{xformers_mask_bw_time:.4f}\",\n",
    "            f\"{calculate_tflops(flops, xformers_mask_bw_time, 10):.2f}\",\n",
    "        ],\n",
    "        [\n",
    "            \"flexattention\",\n",
    "            f\"{flex_ms:.4f}\",\n",
    "            f\"{calculate_tflops(flops, flex_ms, 4):.2f}\",\n",
    "            f\"{flex_bw_ms:.4f}\",\n",
    "            f\"{calculate_tflops(flops, flex_bw_ms, 10):.2f}\",\n",
    "        ],\n",
    "    ]\n",
    "    print(\n",
    "        f\"\\nResults for {score_mod.__name__ if score_mod is not None else mask_mod.__name__}:\"\n",
    "    )\n",
    "    print(\n",
    "        tabulate(\n",
    "            results,\n",
    "            headers=[\n",
    "                \"Operation\",\n",
    "                \"FW Time (ms)\",\n",
    "                \"FW FLOPS (TF/s)\",\n",
    "                \"BW Time (ms)\",\n",
    "                \"BW FLOPS (TF/s)\",\n",
    "            ],\n",
    "            tablefmt=\"grid\",\n",
    "        )\n",
    "    )\n",
    "    if print_mask:\n",
    "        print(f\"\\nBlock Mask:\\n{block_mask}\")\n",
    "\n",
    "    # Clean up to save memory\n",
    "    del query, key, value, gradOut, causal_fa2_out, xformers_out, flex_out\n",
    "    torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Usage\n",
    "\n",
    "Here's a basic example of how to use the FlexAttention API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "output": {
     "id": 512219537823595,
     "loadingStatus": "loaded"
    }
   },
   "outputs": [],
   "source": [
    "def checkerboard(score, batch, head, token_q, token_kv):\n",
    "    score = torch.where(torch.abs(token_kv - token_q) % 1 == 0, score * 0.5, score)\n",
    "    score = torch.where(torch.abs(token_kv - token_q) % 2 == 0, score * 2.0, score)\n",
    "    return score\n",
    "\n",
    "\n",
    "# Create input tensors\n",
    "query = torch.randn(8, 8, 2048, 64, device=\"cuda\", dtype=torch.float32)\n",
    "key = torch.randn(8, 8, 2048, 64, device=\"cuda\", dtype=torch.float32)\n",
    "value = torch.randn(8, 8, 2048, 64, device=\"cuda\", dtype=torch.float32)\n",
    "\n",
    "# Call flex_attention with the checkerboard score modification\n",
    "output = flex_attention(query, key, value, score_mod=checkerboard)\n",
    "\n",
    "# Compile and run\n",
    "compiled_flex_attention = torch.compile(flex_attention)\n",
    "out_compiled = compiled_flex_attention(query, key, value, score_mod=checkerboard)\n",
    "\n",
    "# Check if the results are close\n",
    "torch.testing.assert_close(output, out_compiled, atol=2e-2, rtol=2e-2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Score Modification vs Score Masking\n",
    "We are going to take a brief aside to describe two key concepts that will be important to understand for getting the maximum performance benefits for FlexAttenion.\n",
    "The full api for flex_attention is:\n",
    "```python\n",
    "flex_attention(\n",
    "    query: torch.Tensor,\n",
    "    key: torch.Tensor,\n",
    "    value: torch.Tensor,\n",
    "    score_mod: Optional[Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]] = None,\n",
    "    block_mask: Optional[torch.nn.attention.flex_attention.BlockMask] = None,\n",
    "    scale: Optional[float] = None,\n",
    ")\n",
    "```\n",
    "You may be wondering why we need both a 'score_mod' and a 'block_mask'.\n",
    "1. score_mod functions should be used when you want to mutate score values in the attention weight matrix.\n",
    "2. mask_mod functions should be used when you want to mask scores in the attention weight matrix that are independent of the score value and only rely on positional information.\n",
    "\n",
    "Note: Any block_mask could also be represented with a score_mod, however the performance of the kernel will be suboptimal\n",
    "\n",
    "### Lets walk through causal attention to highlight the differences.\n",
    "\n",
    "The implementation using a score_mod:\n",
    "```Python\n",
    "def causal_bias(score, b, h, q_idx, kv_idx):\n",
    "    return torch.where(q >= kv_idx, score, -float(\"inf\"))\n",
    "```\n",
    "\n",
    "Whenever you are writing a score_mod function that passes through the original score for some elements and sets others to -inf, you should likely be using a mask mod.\n",
    "\n",
    "\n",
    "The implementation using as a mask_mod:\n",
    "```Python\n",
    "The implementation using a mask_mod:\n",
    "def casual_mask(b,h,q_idx, kv_idx):\n",
    "    return q >= kv_idx\n",
    "```\n",
    "As you can see they look very similar, both return scalar tensors. The key differences\n",
    "1. mask_mods return boolean tensors where `True` indicates this score should be calculated, and `False` indicates we that we want to mask out this score\n",
    "2. mask_mods do not take a `score` argument since they are not allowed to depend on actual values during the calculation.\n",
    "\n",
    "\n",
    "### What happens when I use a score_mod + a mask_mod?\n",
    "The score_mod function will be applied to every un-masked element.\n",
    "\n",
    "### I have a mask mod function, how do I create a BlockMask?\n",
    "Great question reader! Besides flex_attention we provide 1 other main API.\n",
    "```python\n",
    "create_block_mask(\n",
    "    mask_mod (Callable): mask_mod function.\n",
    "    B (int): Batch size.\n",
    "    H (int): Number of heads.\n",
    "    Q_LEN (int): Sequence length of query.\n",
    "    KV_LEN (int): Sequence length of key/value.\n",
    "    device (str): Device to run the mask creation on.\n",
    "    KV_BLOCK_SIZE (int): Block size of block mask for each query.\n",
    "    Q_BLOCK_SIZE (int): Block size of block mask for each key/value.\n",
    "    _compile (bool): Whether to compile the mask creation.\n",
    ")\n",
    "```\n",
    "\n",
    "So for the above example the call to flex_attention that would be the most performant is:\n",
    "``` python\n",
    "causal_block_mask = create_block_mask(causal_mask, B, H, M, N)\n",
    "flex_attention(query, key, value, block_mask = causal_block_mask)\n",
    "```\n",
    "B,H,Q_LEN,KV_LEN are the batch_size, num_heads, query_sequence_length, and key_sequence_length.\n",
    "\n",
    "### Why have both?\n",
    "Purely for performance. Causal masking is in fact very sparse. Only the lower triangular portion of the attention scores matter. Without generating a BlockMask we would be doing twice the work needed!\n",
    "Below we will compare the performance difference between the two implementations.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Score Modification Examples\n",
    "\n",
    "Let's explore various score modification examples that can be used with the FlexAttention API. \n",
    "\n",
    "Legend:\n",
    "We are going to be printing a representation of the sparsity found for these score_mod + mask_fns. \n",
    "\n",
    "*  The absence of any block means that it is completely masked out and is not actually needed to compute the final attended output\n",
    "* `██` This block computes full attention between all query and key tokens\n",
    "* `░░` This block is partially masked out, some query tokens attend to some key tokens but some are masked to -inf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Full Attention\n",
    "Applies a \"no-op\" score mod. Leaving the attention scores unchanged."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "output": {
     "id": 1937129110041502,
     "loadingStatus": "loaded"
    }
   },
   "outputs": [],
   "source": [
    "def noop(score, b, h, q_idx, kv_idx):\n",
    "    return score\n",
    "\n",
    "\n",
    "test_mask(noop, print_mask=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Standard Causal Masking\n",
    "Standard Causal Masking is a key technique in autoregressive language models that ensures each token can only attend to itself and previous tokens in the sequence. The block sparse representation shows the lower triangular nature of this mask.\n",
    "\n",
    "See [Score Modification vs Score Masking](#Score-Modification-vs-Score-Masking) for more details on these implementations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "output": {
     "id": 868171238697797,
     "loadingStatus": "loaded"
    }
   },
   "outputs": [],
   "source": [
    "def causal_bias(score, b, h, q_idx, kv_idx):\n",
    "    return torch.where(q_idx >= kv_idx, score, -float(\"inf\"))\n",
    "\n",
    "\n",
    "test_mask(score_mod=causal_bias)\n",
    "\n",
    "\n",
    "def causal_mask(b, h, q_idx, kv_idx):\n",
    "    return q_idx >= kv_idx\n",
    "\n",
    "\n",
    "test_mask(mask_mod=causal_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sliding Window Attention\n",
    "The [Mistral paper](https://arxiv.org/abs/2310.06825) has a very nice visual of this bias and describes it. In essence you define a fixed size \"SLIDING_WINDOW\" and for autogressive decoding you only allow `torch.abs(q_tokens - kv_tokens) < SLIDING_WINDOW` to attend to each other. Typically this is also combined with causal attention. We are going to do this through a a nice pattern, mask composition. Typically masking can can conceptually be done in pieces and then composed together.\n",
    "\n",
    "We are going to write two mask_functions 1 for doing `casual-masking`, and one for doing `windowed-attention` and compose them together to produce the final mask_fn. As we know from earlier, mask_fns return boolean values where a value of `True` indicates that the element should take part in attention.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "output": {
     "id": 1410119323027100,
     "loadingStatus": "loaded"
    }
   },
   "outputs": [],
   "source": [
    "SLIDING_WINDOW = 1024\n",
    "\n",
    "\n",
    "def sliding_window_causal_mask(b, h, q_idx, kv_idx):\n",
    "    causal_mask = q_idx >= kv_idx\n",
    "    windowed_mask = (\n",
    "        q_idx - kv_idx <= SLIDING_WINDOW\n",
    "    )  # We dont need to check the right side of the sliding window since we are applying the causal mask\n",
    "\n",
    "    return causal_mask & windowed_mask\n",
    "\n",
    "\n",
    "test_mask(mask_mod=sliding_window_causal_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prefix LM (Bidirectional + Causal)\n",
    "This T5 achitecture [papers with code](https://paperswithcode.com/method/t5) describes an attention variant that performs prefix attention. Where a certain number of `prefix` tokens are allowed to full attend and then all subsequent tokens perform causal attention. We again compose two mask functions to accomplish this, one for causal masking and one that is based off of the prefix length.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "output": {
     "id": 1917748298667435,
     "loadingStatus": "loaded"
    }
   },
   "outputs": [],
   "source": [
    "PREFIX_LENGTH = 2048\n",
    "\n",
    "\n",
    "def prefix_lm_causal_mask(b, h, q_idx, kv_idx):\n",
    "    prefix_mask = kv_idx <= PREFIX_LENGTH\n",
    "    causal_mask = q_idx >= kv_idx\n",
    "    return prefix_mask | causal_mask\n",
    "\n",
    "\n",
    "test_mask(mask_mod=prefix_lm_causal_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Document Masking\n",
    "Imagine that we have multiple documents of different lengths. We want to mask\n",
    "out the attention between documents, but allow attention between tokens within\n",
    "the same document. We can do this by using a document_id tensor that gives the\n",
    "document that each token belongs to. Then, we can mask out all attention\n",
    "scores where the document_id[q_idx] differs from document_id[kv_idx]\n",
    "\n",
    "\n",
    "Note: We *only* need to compile a new kernel when the `score_mod` changes\n",
    "(it'll automatically detect that using torch.compile infra). This example code\n",
    "is implemented with caching BlockMask, but in general, changing BlockMask\n",
    "*does not* require a recompile.\n",
    "That is, for document masking, we only need to compute a new BlockMask when\n",
    "the document lengths change, *not* a new kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "output": {
     "id": 946844540547075,
     "loadingStatus": "loaded"
    }
   },
   "outputs": [],
   "source": [
    "document_id = torch.zeros(32768, dtype=torch.int, device=\"cuda\")\n",
    "document_id[:4096] = 0\n",
    "document_id[4096:8192] = 1\n",
    "for i in range(8192, 32768, 8192):\n",
    "    document_id[i : i + 8192] = i // 8192 + 1\n",
    "\n",
    "\n",
    "def document_causal_mask(b, h, q_idx, kv_idx):\n",
    "    causal_mask = q_idx >= kv_idx\n",
    "    document_mask = document_id[q_idx] == document_id[kv_idx]\n",
    "    return causal_mask & document_mask\n",
    "\n",
    "\n",
    "test_mask(mask_mod=document_causal_mask, S=32768)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stand-Alone Self-Attention Masking\n",
    "\n",
    "In this case, imagine that we have a 2D image of size (H x W) flattened into a\n",
    "sequence of tokens. We only want to attend to tokens within 8 `pixels`, but\n",
    "from a 2D perspective.\n",
    "\n",
    "We can implement this mask_mod by first translating the 1D position into 2D coordinates. Then, we can simply check if the distance of both coordinates is within the window.\n",
    "\n",
    "For more details check the paper, [Stand-Alone Self-Attention in Vision Models](https://arxiv.org/abs/1906.05909)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "output": {
     "id": 1643933326367268,
     "loadingStatus": "loaded"
    }
   },
   "outputs": [],
   "source": [
    "H = 128\n",
    "W = 128\n",
    "WINDOW = 8\n",
    "\n",
    "\n",
    "def get_x_y(idx):\n",
    "    return idx // W, idx % W\n",
    "\n",
    "\n",
    "def sasa_mask(b, h, q_idx, kv_idx):\n",
    "    q_x, q_y = get_x_y(q_idx)\n",
    "    kv_x, kv_y = get_x_y(kv_idx)\n",
    "    horizontal_mask = (q_x - kv_x).abs() <= WINDOW\n",
    "    vertical_mask = (q_y - kv_y).abs() <= WINDOW\n",
    "    return horizontal_mask & vertical_mask\n",
    "\n",
    "\n",
    "test_mask(mask_mod=sasa_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### NATTEN Masking\n",
    "\n",
    "Consider a 2D image of size (H x W) flattened into a sequence of tokens.\n",
    "Queries attend to keys in a fixed kernel area (K_H x K_W), centered where possible\n",
    "on the query, whilst staying within the canvas and always including the query.\n",
    "\n",
    "This is similar to SASA, except with extra handling to keep the kernel inside the canvas,\n",
    "ensuring that all queries attend to a fixed number of keys.  \n",
    "Keys compare their position to the kernel center, not the query. The kernel center attempts\n",
    "to follow the query position, but is clamped to stay a fixed distance (its half-length) away\n",
    "from the canvas edge.\n",
    "\n",
    "See the [NATTEN repository](https://github.com/SHI-Labs/NATTEN) for more information.  \n",
    "_Note: a more complete implementation of NATTEN would include support for kernel dilation._  \n",
    "_The NATTEN unfused kernel also has features like the ability to cross-attend to register tokens._\n",
    "_This capability is possible to express in Flex Attention but not attempted here._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "H = 128\n",
    "W = 128\n",
    "K_H = 7\n",
    "K_W = 7\n",
    "\n",
    "\n",
    "def get_x_y(idx):\n",
    "    return idx // W, idx % W\n",
    "\n",
    "\n",
    "def natten_mask(\n",
    "    b,\n",
    "    h,\n",
    "    q_idx,\n",
    "    kv_idx,\n",
    "):\n",
    "    q_x, q_y = get_x_y(q_idx)\n",
    "    kv_x, kv_y = get_x_y(kv_idx)\n",
    "    # kernel nominally attempts to center itself on the query, but kernel center\n",
    "    # is clamped to a fixed distance (kernel half-length) from the canvas edge\n",
    "    kernel_x = q_x.clamp(K_W // 2, (W - 1) - K_W // 2)\n",
    "    kernel_y = q_y.clamp(K_H // 2, (H - 1) - K_H // 2)\n",
    "    hori_mask = (kernel_x - kv_x).abs() <= K_W // 2\n",
    "    vert_mask = (kernel_y - kv_y).abs() <= K_H // 2\n",
    "    return hori_mask & vert_mask\n",
    "\n",
    "\n",
    "test_mask(mask_mod=natten_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Alibi Bias\n",
    "The Alibi attention bias was made popular in [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://arxiv.org/abs/2108.12409), and claims to have beneficial properties for length extrapolation at inference \"ALiBi does not add positional embeddings to word embeddings; instead, it biases query-key attention scores with a penalty that is proportional to their distance. \"\n",
    "\n",
    "We are going to implement this 2 ways to highlight a new functionality, the ability to utilize other tensors in your score mod function. Although the function signature doesn't accept other tensors users can implement this via a `closure`. Here we utilize our all too familiar causal mask fn as well as the individual head biases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "output": {
     "id": 515786777538150,
     "loadingStatus": "loaded"
    }
   },
   "outputs": [],
   "source": [
    "# Alibi Bias\n",
    "def generate_alibi_bias():\n",
    "    alibi_bias = []\n",
    "    for h in range(H):\n",
    "        alibi_bias.append(-((h + 1) * 8.0 / H))\n",
    "    alibi_bias = torch.tensor(alibi_bias, device=\"cuda\")\n",
    "    alibi_bias = torch.exp2(alibi_bias)\n",
    "    return alibi_bias\n",
    "\n",
    "\n",
    "alibi_bias = generate_alibi_bias()\n",
    "\n",
    "\n",
    "# In this case we are going to use a mask_mod and a score_mod\n",
    "def causal_mask(b, h, q_idx, kv_idx):\n",
    "    return q_idx >= kv_idx\n",
    "\n",
    "\n",
    "def alibi_and_causal_closure(score, b, h, q_idx, kv_idx):\n",
    "    bias = alibi_bias[h] * (q_idx - kv_idx)\n",
    "    return score + bias\n",
    "\n",
    "\n",
    "def alibi_and_causal_functional(score, b, h, q_idx, kv_idx):\n",
    "    scale = torch.exp2(-((h + 1) * 8.0 / H))\n",
    "    bias = (q_idx - kv_idx) * scale\n",
    "    return score + bias\n",
    "\n",
    "\n",
    "# Correctness check here is simple and only works with mask_fns and not actual score_mods\n",
    "\n",
    "test_mask(\n",
    "    alibi_and_causal_closure,\n",
    "    mask_mod=causal_mask,\n",
    "    skip_correctness=True,\n",
    "    print_mask=False,\n",
    ")\n",
    "test_mask(\n",
    "    alibi_and_causal_functional,\n",
    "    mask_mod=causal_mask,\n",
    "    skip_correctness=True,\n",
    "    print_mask=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tanh Soft-Capping\n",
    "We can also implement tanh soft-capping with this API. Logit softcapping via tanh was popularized in [Gemma 2](https://storage.googleapis.com/deepmind-media/gemma/gemma-2-report.pdf).\n",
    "\n",
    "In this case, there are some nuances. In particular, the standard `tanh`\n",
    "operator in PyTorch (and CUDA/Triton) lowers to a numerically accurate but\n",
    "(relatively) quite slow implementation in SASS. See\n",
    "https://godbolt.org/z/W8afevWv1 for how the SASS looks like.\n",
    "\n",
    "So, in this case, we want to lower the `tanh` into the approximate tanh\n",
    "implementation. We can do so by register a custom operator in PyTorch and then\n",
    "an Inductor lowering."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "output": {
     "id": 856588333051087,
     "loadingStatus": "loaded"
    }
   },
   "outputs": [],
   "source": [
    "# Tanh Soft-Capping\n",
    "@torch.library.custom_op(\"approx::tanh\", mutates_args=())\n",
    "def tanh_approx(inp: torch.Tensor) -> torch.Tensor:\n",
    "    return torch.tanh(inp)\n",
    "\n",
    "\n",
    "@tanh_approx.register_fake\n",
    "def _(inp: torch.Tensor) -> torch.Tensor:\n",
    "    return torch.tanh(inp)\n",
    "\n",
    "\n",
    "from torch._inductor.lowering import make_pointwise, register_lowering\n",
    "\n",
    "# Some internal torch.compile details\n",
    "from torch._inductor.virtualized import ops\n",
    "\n",
    "\n",
    "def tanh_approx_lowering(inp):\n",
    "    fn = partial(ops.inline_asm_elementwise, asm=\"tanh.approx.f32 $0, $1;\")\n",
    "    return make_pointwise(fn)(inp)\n",
    "\n",
    "\n",
    "register_lowering(torch.ops.approx.tanh)(tanh_approx_lowering)\n",
    "\n",
    "\n",
    "class TanhApprox(torch.autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(x):\n",
    "        return torch.ops.approx.tanh(x)\n",
    "\n",
    "    @staticmethod\n",
    "    def setup_context(ctx, inputs, output):\n",
    "        (x,) = inputs\n",
    "        result = output\n",
    "        ctx.save_for_backward(result)\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, grad_output):\n",
    "        (result,) = ctx.saved_tensors\n",
    "        return grad_output * (1 - result * result)\n",
    "\n",
    "\n",
    "tanh_approx = TanhApprox.apply\n",
    "\n",
    "\n",
    "def tanh_soft_cap(score, b, h, q_idx, kv_idx):\n",
    "    score = score / 2\n",
    "    score = tanh_approx(score)\n",
    "    return score * 2\n",
    "\n",
    "\n",
    "# The baseline (xformers) does not have a way to generate tanh-softcapping so we skip correctness checks\n",
    "test_mask(tanh_soft_cap, mask_mod=causal_mask, skip_correctness=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Nested Jagged Tensor\n",
    "\n",
    "Nested Tensors are a tensor subclass that is used to efficiently represent and compute with ragged data. It is possible to utilize FlexAttention with this data to efficiently perform causal attention on batches of sequences with different lengths.\n",
    "\n",
    "Under the hood NJT stores its ragged data as a contiguous data `[[sequence_0], [sequence_1], ..., [Sequence_B]]`, `sum(*),..`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "output": {
     "id": 925610529334684,
     "loadingStatus": "loaded"
    }
   },
   "outputs": [],
   "source": [
    "random.seed(0)\n",
    "torch.manual_seed(0)\n",
    "\n",
    "batch_size = 16\n",
    "n_heads = 16\n",
    "D = 64\n",
    "\n",
    "\n",
    "def prepare_qkv_values(tensor):\n",
    "    return tensor._values.detach().requires_grad_()\n",
    "\n",
    "\n",
    "def build_seq_idx(tensor: torch.Tensor):\n",
    "    offsets = tensor.offsets()\n",
    "    total_length = tensor.offsets()[-1].item()\n",
    "    # Create a range tensor from 0 to total_length\n",
    "    range_tensor = torch.arange(total_length, device=\"cuda\", dtype=torch.int32)\n",
    "\n",
    "    # Use searchsorted to find the index for each position\n",
    "    seq_idx = torch.searchsorted(offsets, range_tensor, right=True) - 1\n",
    "\n",
    "    return seq_idx\n",
    "\n",
    "\n",
    "def create_njt_wrapper(orig_mask_mod, offsets, seq_idx):\n",
    "    \"\"\"Generic Wrapper that converts Dense mask_mod functions to NJT mask_mod functions\"\"\"\n",
    "\n",
    "    def njt_score_mod(b, h, q_idx, kv_idx):\n",
    "        q_nested = q_idx - offsets[seq_idx[q_idx]]\n",
    "        kv_nested = kv_idx - offsets[seq_idx[kv_idx]]\n",
    "        is_same_sequence = seq_idx[q_idx] == seq_idx[kv_idx]\n",
    "        return orig_mask_mod(b, h, q_nested, kv_nested) & is_same_sequence\n",
    "\n",
    "    return njt_score_mod\n",
    "\n",
    "\n",
    "# Dense Score Mod\n",
    "def causal_mask(b, h, q_idx, kv_idx):\n",
    "    return q_idx >= kv_idx\n",
    "    # return torch.where(q_idx >= kv_idx, score, -float(\"inf\"))\n",
    "\n",
    "\n",
    "# Current limitation that the total sequnce length must be divisible by 128\n",
    "sentence_lengths = [random.randint(1, 1024) for _ in range(batch_size - 1)]\n",
    "total = sum(sentence_lengths)\n",
    "sentence_lengths.append(128 - total % 128)\n",
    "total = sum(sentence_lengths)\n",
    "\n",
    "ragged_tensors = [torch.randn(l, n_heads, D, device=\"cuda\") for l in sentence_lengths]\n",
    "query = torch.nested.nested_tensor(\n",
    "    ragged_tensors, layout=torch.jagged, requires_grad=True\n",
    ")\n",
    "key = torch.nested.nested_tensor(\n",
    "    ragged_tensors, layout=torch.jagged, requires_grad=True\n",
    ")\n",
    "value = torch.nested.nested_tensor(\n",
    "    ragged_tensors, layout=torch.jagged, requires_grad=True\n",
    ")\n",
    "\n",
    "# Build the seq_idx lookup table for\n",
    "offsets = query.offsets()\n",
    "seq_idx = build_seq_idx(query)\n",
    "\n",
    "causal_score_mod_njt = create_njt_wrapper(causal_mask, offsets, seq_idx)\n",
    "\n",
    "query_values = prepare_qkv_values(query)\n",
    "key_values = prepare_qkv_values(key)\n",
    "value_values = prepare_qkv_values(value)\n",
    "\n",
    "block_mask = create_block_mask_cached(\n",
    "    causal_score_mod_njt, 1, 1, total, total, device=query_values.device\n",
    ")\n",
    "out_flex = flex_attention(\n",
    "    query_values.view(1, -1, n_heads, D).transpose(1, 2),\n",
    "    key_values.view(1, -1, n_heads, D).transpose(1, 2),\n",
    "    value_values.view(1, -1, n_heads, D).transpose(1, 2),\n",
    "    block_mask=block_mask,\n",
    ")\n",
    "out_sdpa = F.scaled_dot_product_attention(\n",
    "    query.transpose(1, 2),\n",
    "    key.transpose(1, 2),\n",
    "    value.transpose(1, 2),\n",
    "    is_causal=True,\n",
    ")\n",
    "\n",
    "sdpa_outs = []\n",
    "flex_outs = []\n",
    "\n",
    "gradOut = torch.randn_like(out_sdpa)\n",
    "\n",
    "sdpa_outs.append(out_sdpa)\n",
    "out_sdpa.backward(gradOut)\n",
    "sdpa_outs += [query.grad, key.grad, value.grad]\n",
    "\n",
    "flex_outs.append(out_flex)\n",
    "out_flex.backward(gradOut._values.unsqueeze(0))\n",
    "flex_outs += [query_values.grad, key_values.grad, value_values.grad]\n",
    "\n",
    "for flex, sdpa in zip(flex_outs, sdpa_outs):\n",
    "    flex = flex.squeeze(0)\n",
    "    torch.testing.assert_close(flex, sdpa._values, atol=1e-2, rtol=1e-2)\n",
    "\n",
    "\n",
    "print(\"Correctness check passed ✅\")\n",
    "\n",
    "print(block_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Flamingo Cross Attention\n",
    "\n",
    "The 🦩 [Flamingo Paper](https://arxiv.org/pdf/2204.14198) introduced \"a family of visual language models (VLMs)\n",
    "that take as input visual data interleaved with text and produce free-form text as output.\"\n",
    "\n",
    "It utilizes `VisionCrossAttentionMask` to ensure that text only attends to associated images. TorchTune has a good description of this type of masking: [VisionCrossAttentionMask](https://github.com/pytorch/torchtune/blob/bbc48e089b072c7cbaea175bc70501b2193ba482/torchtune/modules/transforms/_transforms.py#L22-L43)\n",
    "\n",
    "This type of attention makes sure that text sequences attend entirely to the preceding image and not to other future/unrelated images.\n",
    "\n",
    "\n",
    "```Python\n",
    "Example:\n",
    "    >>> text = \"<img1><img2>These are two dogs. <img3>This is a cat.\"\n",
    "    >>> image_token_id = 1\n",
    "    >>> tokens = [1, 1, 9673, 527, 1403, 12875, 13, 1, 1115, 374, 264, 8415]\n",
    "    >>> transform = VisionCrossAttentionMask(tile_size=400, patch_size=40, image_token_id=1)\n",
    "    >>> intervals = transform._get_image_attention_intervals(tokens)\n",
    "    >>> print(intervals)\n",
    "    [[0, 7], [1, 7], [7, 12]]\n",
    "```\n",
    "\n",
    "In the above case, we would generate a mask of\n",
    "12 x sum(image_tokens_1 + image_tokens_2 + image_tokens_3)\n",
    "\n",
    "assuming image_tokens are size 3\n",
    "```\n",
    "\n",
    "          img1    img2   img3\n",
    "     1   █ █ █ | ░ ░ ░ | ░ ░ ░\n",
    "     1   █ █ █ | █ █ █ | ░ ░ ░\n",
    "  9673   █ █ █ | █ █ █ | ░ ░ ░\n",
    "   527   █ █ █ | █ █ █ | ░ ░ ░\n",
    "  1403   █ █ █ | █ █ █ | ░ ░ ░\n",
    " 12875   █ █ █ | █ █ █ | ░ ░ ░\n",
    "    13   █ █ █ | █ █ █ | ░ ░ ░\n",
    "     1   ░ ░ ░ | ░ ░ ░ | █ █ █\n",
    "  1115   ░ ░ ░ | ░ ░ ░ | █ █ █\n",
    "   374   ░ ░ ░ | ░ ░ ░ | █ █ █\n",
    "   264   ░ ░ ░ | ░ ░ ░ | █ █ █\n",
    "  8415   ░ ░ ░ | ░ ░ ░ | █ █ █\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "output": {
     "id": 1526626144609389,
     "loadingStatus": "loaded"
    }
   },
   "outputs": [],
   "source": [
    "# Given information\n",
    "num_tokens = 12\n",
    "num_images = 3\n",
    "image_token_length = 3\n",
    "num_image_tokens = num_images * image_token_length\n",
    "intervals = [[0, 7], [1, 7], [7, 12]]\n",
    "# This is only needed if your images have different number of tokens per image\n",
    "# If they are all the same number of tokens you can use image_idx = kv_idx // image_token_length\n",
    "image_boundaries = [image_token_length * i for i in range(num_images)]\n",
    "image_boundaries = (\n",
    "    [0] * image_token_length + [1] * image_token_length + [2] * image_token_length\n",
    ")\n",
    "\n",
    "image_boundaries = torch.tensor(image_boundaries, dtype=torch.long, device=\"cuda\")\n",
    "intervals = torch.tensor(intervals, dtype=torch.long, device=\"cuda\")\n",
    "\n",
    "\n",
    "def vision_x_attention_mask(b, h, q_idx, kv_idx):\n",
    "    image_idx = image_boundaries[kv_idx]\n",
    "    interval = intervals[image_idx]\n",
    "    return (q_idx >= interval[0]) & (q_idx < interval[1])\n",
    "\n",
    "\n",
    "mask = create_mask(vision_x_attention_mask, 1, 1, num_tokens, num_image_tokens, \"cuda\")\n",
    "\n",
    "print(mask)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dev",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}