{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SDPA Prefill with Paged Attention\n",
    "\n",
    "This notebook illustrates how the cuDNN's frontend scaled dot product attention operator can be used with paged K/V caches, specifically for prefill. For a simpler introduction to the scaled dot product attention operator, please refer to [samples/python/50_scaled_dot_product_attention.ipynb](https://github.com/NVIDIA/cudnn-frontend/blob/main/samples/python/50_scaled_dot_product_attention.ipynb)\n",
    "\n",
    "The full documentation of cuDNN's scaled dot production attention operator can be found in: [docs/operations/Attention.md#scaled-dot-product-attention](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Attention.md#scaled-dot-product-attention). The python test code for the full set of features can be found in: [test/python/test_mhas.py](https://github.com/NVIDIA/cudnn-frontend/blob/main/test/python/test_mhas.py)\n",
    "\n",
    "More details on paged attention can be found in the [PagedAttention paper](https://arxiv.org/abs/2309.06180).\n",
    "\n",
    "This notebook specifically illustrates the following:\n",
    "- SDPA Prefill\n",
    "- Variable sequence lengths\n",
    "- Q-tensor in a dense format\n",
    "- Paged Attention\n",
    "- Running the same graph with variable sequence lengths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NVIDIA/cudnn-frontend/blob/main/samples/python/52_scaled_dot_product_attention.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Prerequisites and Setup\n",
    "This notebook requires an NVIDIA GPU A100 or newer. If running on Colab, go to Runtime → Change runtime type → Hardware accelerator and select a GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get_ipython().system('nvidia-smi')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get_ipython().system('pip install nvidia-cudnn-cu12')\n",
    "# get_ipython().system('pip install nvidia-cudnn-frontend')\n",
    "# get_ipython().system('pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cudnn\n",
    "import torch\n",
    "import math\n",
    "\n",
    "torch.manual_seed(42)\n",
    "handle = cudnn.create_handle()\n",
    "\n",
    "assert torch.cuda.is_available()\n",
    "assert (\n",
    "    torch.cuda.get_device_capability()[0] >= 8\n",
    "), \"SDPA operation is only supported on SM80 architecture (Ampere) or above\"\n",
    "\n",
    "assert (\n",
    "    cudnn.backend_version() >= 90500\n",
    "), \"SDPA operation is only supported cuDNN version 9.5.0 or above\"\n",
    "\n",
    "# An issue is preventing packed Q-tensors in versions prior to 9.10.0\n",
    "packed_Q = cudnn.backend_version() >= 91000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Problem sizes and Q/K/V setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the query, key, value, and output GPU tensors using PyTorch. However, the user may use any DLPack compatible tensor instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "b = 4  # batch size\n",
    "h_q = 12  # query number of heads\n",
    "h_kv = 12  # key and value number of heads\n",
    "s_q = 128  # maximum sequence length for Q\n",
    "s_kv = 128  # maximum sequence length for K/V\n",
    "d = 64  # embedding dimension per head\n",
    "\n",
    "block_size_k = block_size_v = (\n",
    "    8  # block size to be used by the non contiguous K/V containers\n",
    ")\n",
    "attn_scale = 1.0 / math.sqrt(d)\n",
    "\n",
    "# BHSD (batch, sequence_length, num_head, dims_per_head) logcial tensor layout\n",
    "dims_q = (b, h_q, s_q, d)\n",
    "dims_kv = (b, h_kv, s_kv, d)\n",
    "# BSHD physical tensor layout (this is required for packed Q-tensors)\n",
    "strides_q = (s_q * h_q * d, d, h_q * d, 1)\n",
    "strides_kv = (s_kv * h_kv * d, d, h_kv * d, 1)\n",
    "\n",
    "# Randomly initialize the query, key, and value tensors.\n",
    "q_gpu = torch.randn(b * s_q * h_q * d).half().cuda().as_strided(dims_q, strides_q)\n",
    "k_gpu = torch.randn(b * s_kv * h_kv * d).half().cuda().as_strided(dims_kv, strides_kv)\n",
    "v_gpu = torch.randn(b * s_kv * h_kv * d).half().cuda().as_strided(dims_kv, strides_kv)\n",
    "o_gpu = torch.empty(b * s_q * h_q * d).half().cuda().as_strided(dims_q, strides_q)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Setup actual sequence lengths and ragged offsets\n",
    "While we defined `s_q` as the maximum sequence length for Q and `s_kv` as the maximum sequence length for K/V, not all sequences have the same length. Therefore we specify actual sequence lengths in this section.\n",
    "\n",
    "While optional for Q, actual sequence lengths are required when using paged K/V caches. To keep things simple, we set initialize the actual sequence lengths for KV to the maximum sequence length `s` in this example, but we will specify random sequence lengths for Q.\n",
    "\n",
    "Lastly, when Q is in a packed format, we also need to create a ragged offset tensor. This is a tensor that indicates the start of each sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @brief Helper function to return variable sequence lengths, along with a ragged offset tensor inidicating the start of each sequence\n",
    "def create_variable_seq_lens(b, s):\n",
    "    seq_len_q_gpu = torch.randint(1, s, (b, 1, 1, 1), device=\"cuda\")\n",
    "\n",
    "    # Create a [b+1, 1, 1, 1] ragged offset tensor\n",
    "    q_ragged_offset_gpu = (\n",
    "        torch.cat(\n",
    "            (\n",
    "                torch.zeros(1, 1, 1, 1, dtype=torch.int32, device=\"cuda\"),\n",
    "                torch.cumsum(seq_len_q_gpu, dim=0),\n",
    "            )\n",
    "        )\n",
    "        * h_q\n",
    "        * d\n",
    "    )\n",
    "    return seq_len_q_gpu, q_ragged_offset_gpu\n",
    "\n",
    "\n",
    "# For Q, randomly generate sequence lengths between [1,s)\n",
    "seq_len_q_gpu, q_ragged_offset_gpu = create_variable_seq_lens(b, s_q)\n",
    "\n",
    "# For KV, set to s for all batches, just to keep this notebook sample simple\n",
    "seq_len_kv_gpu = torch.full((b, 1, 1, 1), s_kv, device=\"cuda\")\n",
    "\n",
    "print(seq_len_q_gpu)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####  Generate containers and block tables for K and V\n",
    "\n",
    "In a real world scenario, container and block table tensors are generated by other parts of the model. For illustration purposes in this example, we provide a helper function to generate a trivial container from contiguous K and V caches. \n",
    "The helper function basically takes e.g., the K-cache and splits up the sequence (`S`) dimension in different blocks of length `block_size`. The resulting block table then helps identify which block belongs to which sequence ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @brief Helper function to create a non contiguous container in blocks of block_size from a contiguous tensor\n",
    "def create_container_and_block_table(tensor, block_size):\n",
    "    B, H, S, D = tensor.shape\n",
    "    # num_blocks = math.ceil(S/block_size) * B\n",
    "    blocks_per_batch = math.ceil(S / block_size)\n",
    "\n",
    "    # Only needed if S is not a multiple of block_size\n",
    "    padding_seq = (blocks_per_batch * block_size) - S\n",
    "    if padding_seq > 0:\n",
    "        zeros = torch.zeros(B, H, padding_seq, D, device=\"cuda\", dtype=tensor.dtype)\n",
    "        cat_tensor = torch.cat((tensor, zeros), axis=2)\n",
    "    else:\n",
    "        cat_tensor = tensor\n",
    "\n",
    "    # Create a container by splitting on the S dimension and concatenating at the block dimension\n",
    "    # Its dimensions are [num_blocks, H, block_size, D] with num_blocks = B * blocks_per_batch\n",
    "    container = torch.cat((cat_tensor.clone()).chunk(blocks_per_batch, dim=2), dim=0)\n",
    "\n",
    "    # Create the block table\n",
    "    table_size = math.ceil(S / block_size)\n",
    "    block_table_temp = torch.linspace(\n",
    "        0, B * table_size - 1, B * table_size, device=\"cuda\", dtype=torch.int32\n",
    "    ).reshape(table_size, 1, B, 1)\n",
    "    block_table_temp = torch.transpose(block_table_temp, 0, 2)\n",
    "\n",
    "    # Make batch size outer dimension (cuDNN backend preference)\n",
    "    block_table = (\n",
    "        torch.zeros(blocks_per_batch * B)\n",
    "        .int()\n",
    "        .cuda()\n",
    "        .as_strided(\n",
    "            (B, 1, blocks_per_batch, 1), (blocks_per_batch, blocks_per_batch, 1, 1)\n",
    "        )\n",
    "    )\n",
    "    block_table.copy_(block_table_temp)\n",
    "\n",
    "    return (container, block_table)\n",
    "\n",
    "\n",
    "# Create non contiguous containers with block tables for K and V from the contiguous k_gpu and v_gpu\n",
    "container_k_gpu, block_table_k_gpu = create_container_and_block_table(\n",
    "    k_gpu, block_size_k\n",
    ")\n",
    "container_v_gpu, block_table_v_gpu = create_container_and_block_table(\n",
    "    v_gpu, block_size_v\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Graph creation\n",
    "\n",
    "Create the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = cudnn.pygraph(\n",
    "    io_data_type=cudnn.data_type.HALF,\n",
    "    intermediate_data_type=cudnn.data_type.FLOAT,\n",
    "    compute_data_type=cudnn.data_type.FLOAT,\n",
    ")\n",
    "\n",
    "q = graph.tensor_like(q_gpu)\n",
    "\n",
    "if packed_Q:\n",
    "    q_ragged_offset = graph.tensor_like(q_ragged_offset_gpu)\n",
    "    q.set_ragged_offset(\n",
    "        q_ragged_offset\n",
    "    )  # With Q in a packed layout, we need to indicate the ragged offset\n",
    "\n",
    "container_k = graph.tensor_like(container_k_gpu)\n",
    "container_v = graph.tensor_like(container_v_gpu)\n",
    "block_table_k = graph.tensor_like(block_table_k_gpu)\n",
    "block_table_v = graph.tensor_like(block_table_v_gpu)\n",
    "\n",
    "seq_len_q = graph.tensor_like(seq_len_q_gpu)\n",
    "seq_len_kv = graph.tensor_like(seq_len_kv_gpu)\n",
    "\n",
    "o, _ = graph.sdpa(\n",
    "    name=\"sdpa\",\n",
    "    q=q,\n",
    "    k=container_k,  # Container K: non contiguous container with K blocks\n",
    "    v=container_v,  # Container V: non contiguous container with V blocks\n",
    "    is_inference=False,\n",
    "    attn_scale=attn_scale,\n",
    "    use_causal_mask=True,\n",
    "    use_padding_mask=True,\n",
    "    seq_len_q=seq_len_q,\n",
    "    seq_len_kv=seq_len_kv,\n",
    "    paged_attention_k_table=block_table_k,  # Block Table K: Tensor containing offsets to the container with K blocks\n",
    "    paged_attention_v_table=block_table_v,  # Block Table V: Tensor containing offsets to the container with V blocks\n",
    "    paged_attention_max_seq_len_kv=s_kv,  # The maximum sequence length for K caches (this is optional, but recommended)\n",
    ")\n",
    "\n",
    "o.set_output(True).set_dim(dims_q).set_stride(strides_q)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph.validate()\n",
    "graph.build_operation_graph()\n",
    "graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])\n",
    "graph.check_support()\n",
    "graph.build_plans()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Graph Execution\n",
    "Execute the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "variant_pack = {\n",
    "    q: q_gpu,\n",
    "    container_k: container_k_gpu,\n",
    "    container_v: container_v_gpu,\n",
    "    block_table_k: block_table_k_gpu,\n",
    "    block_table_v: block_table_v_gpu,\n",
    "    seq_len_q: seq_len_q_gpu,\n",
    "    seq_len_kv: seq_len_kv_gpu,\n",
    "    o: o_gpu,\n",
    "}\n",
    "if packed_Q:\n",
    "    variant_pack[q_ragged_offset] = q_ragged_offset_gpu\n",
    "\n",
    "print(\"First execution\")\n",
    "workspace = torch.empty(graph.get_workspace_size(), device=\"cuda\", dtype=torch.uint8)\n",
    "graph.execute(variant_pack, workspace)\n",
    "torch.cuda.synchronize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run the PyTorch reference and compare against cuDNN's output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compare_against_torch_ref(q_gpu, k_gpu, v_gpu, o_gpu, seq_len_q_gpu):\n",
    "    q_gpu_packed = q_gpu.detach().float().requires_grad_()\n",
    "    k_ref = k_gpu.detach().float().requires_grad_()\n",
    "    v_ref = v_gpu.detach().float().requires_grad_()\n",
    "\n",
    "    (b, h, s_q, d) = q_gpu_packed.shape\n",
    "    s_kv = k_gpu.shape[2]\n",
    "\n",
    "    mask = torch.ones(b, s_q, s_kv, dtype=torch.bool, device=\"cuda\")\n",
    "\n",
    "    # Create attention mask for variable lengths in Q\n",
    "    mask = torch.ones(b, s_q, s_kv, dtype=torch.bool, device=\"cuda\")\n",
    "    for i, length in enumerate(seq_len_q_gpu):\n",
    "        mask[i, length:, :] = False\n",
    "\n",
    "    # Create attention mask for variable lengths in KV\n",
    "    for i, length in enumerate(seq_len_kv_gpu):\n",
    "        mask[i, :, length:] = False\n",
    "    # Causal masking\n",
    "    for i in range(s_q):\n",
    "        mask[:, i, i + 1 :] = False\n",
    "\n",
    "    # Expand mask to match attention shape\n",
    "    mask = mask.unsqueeze(1)\n",
    "\n",
    "    o_ref = None\n",
    "    if packed_Q:\n",
    "        # Create unpacked tensor with proper shape\n",
    "        # Convert bhsd to bshd logical layout and flatten\n",
    "        uniform_tensor = torch.zeros(b, s_q, h, d).to(\n",
    "            dtype=q_gpu_packed.dtype, device=q_gpu_packed.device\n",
    "        )\n",
    "        q_gpu_packed_thd = torch.einsum(\"bhsd->bshd\", q_gpu_packed).reshape(\n",
    "            b * s_q, h, d\n",
    "        )\n",
    "\n",
    "        # Copy the data from the packed tensor to the unpacked tensor\n",
    "        start_idx = 0\n",
    "        for i in range(b):\n",
    "            s = seq_len_q_gpu[i]\n",
    "            uniform_tensor[i, 0:s, :, :] = q_gpu_packed_thd[\n",
    "                start_idx : start_idx + s, :, :\n",
    "            ]\n",
    "            start_idx += s\n",
    "\n",
    "        # Convert back to bhsd logical layout\n",
    "        q_unpacked_ref = torch.einsum(\"bshd->bhsd\", uniform_tensor)\n",
    "\n",
    "        o_ref = torch.nn.functional.scaled_dot_product_attention(\n",
    "            q_unpacked_ref,\n",
    "            k_ref,\n",
    "            v_ref,\n",
    "            is_causal=False,\n",
    "            scale=attn_scale,\n",
    "            attn_mask=mask,\n",
    "        )\n",
    "    else:\n",
    "        o_ref = torch.nn.functional.scaled_dot_product_attention(\n",
    "            q_gpu_packed,\n",
    "            k_ref,\n",
    "            v_ref,\n",
    "            is_causal=False,\n",
    "            scale=attn_scale,\n",
    "            attn_mask=mask,\n",
    "        )\n",
    "\n",
    "    torch.testing.assert_close(o_ref, o_gpu.float(), atol=5e-3, rtol=3e-3)\n",
    "\n",
    "\n",
    "compare_against_torch_ref(q_gpu, k_gpu, v_gpu, o_gpu, seq_len_q_gpu)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NVIDIA/cudnn-frontend/blob/main/samples/python/52_scaled_dot_product_attention.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Executing the same graph with different sequence lengths\n",
    "\n",
    "Note that the graph construction we went through earlier is a one-time cost. We can reuse the same graph for different actual sequence lengths. We illustrate this below by creating new variable sequence lengths for Q."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seq_len_q_gpu, q_ragged_offset_gpu = create_variable_seq_lens(b, s_q)\n",
    "print(seq_len_q_gpu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "variant_pack = {\n",
    "    q: q_gpu,\n",
    "    container_k: container_k_gpu,\n",
    "    container_v: container_v_gpu,\n",
    "    block_table_k: block_table_k_gpu,\n",
    "    block_table_v: block_table_v_gpu,\n",
    "    seq_len_q: seq_len_q_gpu,\n",
    "    seq_len_kv: seq_len_kv_gpu,\n",
    "    o: o_gpu,\n",
    "}\n",
    "\n",
    "if packed_Q:\n",
    "    variant_pack[q_ragged_offset] = q_ragged_offset_gpu\n",
    "\n",
    "print(\"Second execution\")\n",
    "workspace = torch.empty(graph.get_workspace_size(), device=\"cuda\", dtype=torch.uint8)\n",
    "graph.execute(variant_pack, workspace)\n",
    "torch.cuda.synchronize()\n",
    "compare_against_torch_ref(q_gpu, k_gpu, v_gpu, o_gpu, seq_len_q_gpu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cudnn.destroy_handle(handle)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
