### Install PoPE-pytorch Source: https://github.com/hunterhogan/pope-pytorch/blob/main/README.md Install the library using pip. ```shell pip install PoPE-pytorch ``` -------------------------------- ### Initialize and Use PoPE Source: https://github.com/hunterhogan/pope-pytorch/blob/main/README.md Initialize PoPE with dimension and heads, then apply to queries and keys for attention. ```python import torch from PoPE_pytorch import PoPE # define pope pope = PoPE(64, heads = 8) # pass in sequence length pos_emb = pope(1024) # queries and keys in attention q = torch.randn(1, 8, 1024, 64) k = torch.randn(1, 8, 1024, 64) # training rotated_q, rotated_k = pope.apply_pope_to_qk(pos_emb, q, k) # inference rotated_q, rotated_k = pope.apply_pope_to_qk(pos_emb, q[..., -1:, :], k) ``` -------------------------------- ### Initialize and Use AxialPoPE for Video Source: https://github.com/hunterhogan/pope-pytorch/blob/main/README.md Initialize AxialPoPE for video data, splitting the dimension across time and spatial dimensions. ```python # for video (e.g. 8 frames, 16x16 frames) # split 96 dim into 32 (t), 32 (x), 32 (y) pope_video = AxialPoPE( dim = 96, heads = 8, axial_dims = (32, 32, 32) ) pos_emb_video = pope_video((8, 16, 16)) # (2048, 96) frequencies # queries and keys # then apply to q, k as usual q = torch.randn(1, 8, 2048, 96) k = torch.randn(1, 8, 2048, 96) rotated_q, rotated_k = AxialPoPE.apply_pope_to_qk(pos_emb_video, q, k) ``` -------------------------------- ### Initialize and Use AxialPoPE for Images Source: https://github.com/hunterhogan/pope-pytorch/blob/main/README.md Initialize AxialPoPE for image data, splitting the dimension across specified axial dimensions. ```python import torch from PoPE_pytorch import AxialPoPE # axial pope for images (e.g. 16x16) # split 64 dim into 32 (x) and 32 (y) pope = AxialPoPE( dim = 64, heads = 8, axial_dims = (32, 32) ) pos_emb = pope((16, 16)) # (256, 64) frequencies ``` -------------------------------- ### Fused Flash Attention with PoPE Source: https://github.com/hunterhogan/pope-pytorch/blob/main/README.md Utilize PoPE within a fused Flash Attention implementation for efficiency. ```python import torch from PoPE_pytorch import PoPE, flash_attn_with_pope # pope pope = PoPE(dim = 32, heads = 8).cuda() # queries, keys, values for attention q = torch.randn(2, 8, 1024, 64).cuda() k = torch.randn(2, 8, 1024, 64).cuda() v = torch.randn(2, 8, 1024, 64).cuda() pos_emb = pope(1024) mask = torch.ones((2, 1024)).bool().cuda() out = flash_attn_with_pope(q, k, v, pos_emb = pos_emb, causal = True, mask = mask) assert out.shape == (2, 8, 1024, 64) ``` -------------------------------- ### Instantiate and Apply PoPE for 1D Sequences Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt Instantiate the PoPE module and generate positional embeddings for a given sequence length. Apply these embeddings to query and key tensors for training or inference. ```python import torch from PoPE_pytorch import PoPE # --- Instantiate --- # dim: per-head rotation dimension (must be <= head_dim in Q/K) # heads: number of attention heads pope = PoPE(dim=64, heads=8, theta=10000) # --- Generate positional embeddings for a sequence of length 1024 --- pos_emb = pope(1024) # pos_emb.freqs shape: (1024, 64) # pos_emb.bias shape: (8, 64) # --- Apply to Q and K (training: full context) --- q = torch.randn(1, 8, 1024, 64) k = torch.randn(1, 8, 1024, 64) rotated_q, rotated_k = pope.apply_pope_to_qk(pos_emb, q, k) assert rotated_q.shape == (1, 8, 1024, 128) # dim doubles after polar interleaving assert rotated_k.shape == (1, 8, 1024, 128) # --- Apply to Q and K (inference: single new token, cached keys) --- rotated_q_inf, rotated_k_inf = pope.apply_pope_to_qk(pos_emb, q[..., -1:, :], k) assert rotated_q_inf.shape == (1, 8, 1, 128) assert rotated_k_inf.shape == (1, 8, 1024, 128) # --- Use explicit position tensor (e.g. for packed/offset sequences) --- positions = torch.arange(512, 512 + 1024).float() pos_emb_offset = pope(positions) ``` -------------------------------- ### Instantiate and Apply AxialPoPE for Multi-Dimensional Grids Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt Instantiate AxialPoPE for multi-dimensional grids like images or video. Generate positional embeddings using grid dimensions and apply them to query and key tensors. ```python import torch from PoPE_pytorch import AxialPoPE # --- 2D image grid (16×16 patches, head_dim=64 split as 32+32) --- pope_2d = AxialPoPE(dim=64, heads=8, axial_dims=(32, 32)) pos_emb_2d = pope_2d((16, 16)) # auto meshgrid → (256, 64) freqs assert pos_emb_2d.freqs.shape == (256, 64) q = torch.randn(1, 8, 256, 64) k = torch.randn(1, 8, 256, 64) rotated_q, rotated_k = AxialPoPE.apply_pope_to_qk(pos_emb_2d, q, k) assert rotated_q.shape == (1, 8, 256, 128) ``` -------------------------------- ### 3D Positional Embeddings with AxialPoPE Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt Demonstrates the creation and application of 3D positional embeddings using AxialPoPE for video data. Shows both automatic meshgrid generation and manual position definition. ```python pope_3d = AxialPoPE(dim=96, heads=8, axial_dims=(32, 32, 32)) pos_emb_3d = pope_3d((8, 16, 16))) # auto meshgrid → (2048, 96) freqs assert pos_emb_3d.freqs.shape == (2048, 96) q3 = torch.randn(1, 8, 2048, 96) k3 = torch.randn(1, 8, 2048, 96) rotated_q3, rotated_k3 = AxialPoPE.apply_pope_to_qk(pos_emb_3d, q3, k3) assert rotated_q3.shape == (1, 8, 2048, 192) ``` ```python # Manual positions (e.g. irregular or cropped grids) \ custom_pos = AxialPoPE.get_grid_positions(4, 4, device='cpu') # (16, 2) \ pos_emb_manual = pope_2d(custom_pos) ``` -------------------------------- ### Apply PoPE to QK Tensors with Customizations Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt Use the standalone `apply_pope_to_qk` function to apply positional embeddings. Supports partial rotation, custom magnitude activations, and complex tensor output. ```python import torch import torch.nn.functional as F from PoPE_pytorch import PoPE from PoPE_pytorch.pope import apply_pope_to_qk pope = PoPE(dim=32, heads=4) # only rotate the first 32 of 64 dims (partial) pos_emb = pope(256) q = torch.randn(2, 4, 256, 64) k = torch.randn(2, 4, 256, 64) # Partial rotation: first 32 dims rotated, remaining 32 passed through unchanged rotated_q, rotated_k = apply_pope_to_qk(pos_emb, q, k) assert rotated_q.shape == (2, 4, 256, 96) # 32*2 + 32 = 96 # Using ReLU instead of softplus as magnitude activation rotated_q2, rotated_k2 = apply_pope_to_qk(pos_emb, q, k, to_magnitude=F.relu) # Return complex tensors instead of interleaved real rotated_q_c, rotated_k_c = apply_pope_to_qk(pos_emb, q, k, return_complex=True) assert rotated_q_c.is_complex() ``` -------------------------------- ### Flash Attention with PoPE Integration Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt A full attention function that fuses PoPE rotations into a Triton-based flash attention kernel or falls back to scaled_dot_product_attention. Supports various features like causal masking, GQA, and different head layouts. ```python import torch from PoPE_pytorch import PoPE, flash_attn_with_pope pope = PoPE(dim=32, heads=8).cuda() pos_emb = pope(1024) # --- Basic causal self-attention (training) --- q = torch.randn(2, 8, 1024, 64).cuda() k = torch.randn(2, 8, 1024, 64).cuda() v = torch.randn(2, 8, 1024, 64).cuda() mask = torch.ones(2, 1024, dtype=torch.bool).cuda() out = flash_attn_with_pope(q, k, v, pos_emb=pos_emb, causal=True, mask=mask) assert out.shape == (2, 8, 1024, 64) ``` ```python # --- Inference: single new token, cached context (causal) --- q_new = torch.randn(2, 8, 1, 64).cuda() out_inf = flash_attn_with_pope(q_new, k, v, pos_emb=pos_emb, causal=True) assert out_inf.shape == (2, 8, 1, 64) ``` ```python # --- Heads-last layout (b, n, h, d) --- q_hl = torch.randn(2, 1024, 8, 64).cuda() k_hl = torch.randn(2, 1024, 8, 64).cuda() v_hl = torch.randn(2, 1024, 8, 64).cuda() out_hl = flash_attn_with_pope(q_hl, k_hl, v_hl, pos_emb=pos_emb, causal=False, head_dimension_at_first=False) assert out_hl.shape == (2, 1024, 8, 64) ``` ```python # --- Force non-fused (CPU or debugging) --- out_manual = flash_attn_with_pope(q, k, v, pos_emb=pos_emb, causal=True, fused=False) ``` ```python # --- Custom softmax scale + dropout --- out_scaled = flash_attn_with_pope(q, k, v, pos_emb=pos_emb, causal=True, softmax_scale=0.125, dropout=0.1) ``` -------------------------------- ### flash_attn_with_pope Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt A complete multi-head self/cross-attention function that fuses PoPE rotations into a Triton-based flash attention kernel or falls back to `torch.nn.functional.scaled_dot_product_attention`. Supports various masking, dropout, GQA, and layouts. ```APIDOC ## flash_attn_with_pope ### Description A complete multi-head self/cross-attention function that fuses PoPE rotations into a Triton-based flash attention kernel (when available on CUDA) or falls back to `torch.nn.functional.scaled_dot_product_attention`. Supports causal masking, key-padding masks, dropout, grouped query attention, and both heads-first `(b, h, n, d)` and heads-last `(b, n, h, d)` layouts. ### Method `flash_attn_with_pope(q, k, v, pos_emb, causal=False, mask=None, softmax_scale=None, dropout=0.0, fused=True, head_dimension_at_first=True)` ### Parameters - **q** (torch.Tensor) - Query tensor. - **k** (torch.Tensor) - Key tensor. - **v** (torch.Tensor) - Value tensor. - **pos_emb** (torch.Tensor) - Positional embedding tensor. - **causal** (bool, optional) - Whether to apply causal masking. Defaults to False. - **mask** (torch.Tensor, optional) - Key-padding mask. Defaults to None. - **softmax_scale** (float, optional) - Softmax scale factor. Defaults to None. - **dropout** (float, optional) - Dropout rate for attention weights. Defaults to 0.0. - **fused** (bool, optional) - Whether to use the fused kernel. Defaults to True. - **head_dimension_at_first** (bool, optional) - Specifies if the head dimension is the first dimension after batch. Defaults to True. ### Request Example ```python import torch from PoPE_pytorch import PoPE, flash_attn_with_pope pope = PoPE(dim=32, heads=8).cuda() pos_emb = pope(1024) # --- Basic causal self-attention (training) --- q = torch.randn(2, 8, 1024, 64).cuda() k = torch.randn(2, 8, 1024, 64).cuda() v = torch.randn(2, 8, 1024, 64).cuda() mask = torch.ones(2, 1024, dtype=torch.bool).cuda() out = flash_attn_with_pope(q, k, v, pos_emb=pos_emb, causal=True, mask=mask) assert out.shape == (2, 8, 1024, 64) # --- Inference: single new token, cached context (causal) --- q_new = torch.randn(2, 8, 1, 64).cuda() out_inf = flash_attn_with_pope(q_new, k, v, pos_emb=pos_emb, causal=True) assert out_inf.shape == (2, 8, 1, 64) # --- Heads-last layout (b, n, h, d) --- q_hl = torch.randn(2, 1024, 8, 64).cuda() k_hl = torch.randn(2, 1024, 8, 64).cuda() v_hl = torch.randn(2, 1024, 8, 64).cuda() out_hl = flash_attn_with_pope(q_hl, k_hl, v_hl, pos_emb=pos_emb, causal=False, head_dimension_at_first=False) assert out_hl.shape == (2, 1024, 8, 64) # --- Force non-fused (CPU or debugging) --- out_manual = flash_attn_with_pope(q, k, v, pos_emb=pos_emb, causal=True, fused=False) # --- Custom softmax scale + dropout --- out_scaled = flash_attn_with_pope(q, k, v, pos_emb=pos_emb, causal=True, softmax_scale=0.125, dropout=0.1) ``` ### Response #### Success Response (200) - **out** (torch.Tensor) - The output tensor after applying attention. ``` -------------------------------- ### AxialPoPE Module Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt `AxialPoPE` extends PoPE to multi-dimensional grids by splitting the feature dimension across independent axial embeddings. The `forward` method accepts position tensors or grid sizes to generate embeddings for multi-dimensional data. ```APIDOC ## AxialPoPE — Multi-Dimensional Positional Embedding ### Description `AxialPoPE` extends PoPE to multi-dimensional grids by splitting the feature dimension across independent axial embeddings (e.g. x and y for images, t/x/y for video). The `forward` method accepts either an explicit `(N, num_axes)` position tensor or a tuple of grid sizes, which auto-generates a meshgrid. ### Usage (2D Image Grid) ```python import torch from PoPE_pytorch import AxialPoPE # --- 2D image grid (16×16 patches, head_dim=64 split as 32+32) --- pope_2d = AxialPoPE(dim=64, heads=8, axial_dims=(32, 32)) pos_emb_2d = pope_2d((16, 16)) # auto meshgrid → (256, 64) freqs assert pos_emb_2d.freqs.shape == (256, 64) q = torch.randn(1, 8, 256, 64) k = torch.randn(1, 8, 256, 64) rotated_q, rotated_k = AxialPoPE.apply_pope_to_qk(pos_emb_2d, q, k) assert rotated_q.shape == (1, 8, 256, 128) ``` ``` -------------------------------- ### apply_pope_to_qk Function Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt A standalone function (also a static method on `PoPE` and `AxialPoPE`) that applies a `PolarEmbedReturn` to query and key tensors. It supports partial rotation, complex-number output, and configurable magnitude activations. ```APIDOC ## apply_pope_to_qk — Standalone Function ### Description `apply_pope_to_qk` is a standalone `@autocast`-safe function (also available as a static method on `PoPE` and `AxialPoPE`) that applies a `PolarEmbedReturn` to query and key tensors. It supports partial rotation (when `rotate_dim < head_dim`), complex-number output mode, and a configurable magnitude activation (default: `F.softplus`). ### Usage ```python import torch import torch.nn.functional as F from PoPE_pytorch import PoPE from PoPE_pytorch.pope import apply_pope_to_qk pope = PoPE(dim=32, heads=4) # only rotate the first 32 of 64 dims (partial) pos_emb = pope(256) q = torch.randn(2, 4, 256, 64) k = torch.randn(2, 4, 256, 64) # Partial rotation: first 32 dims rotated, remaining 32 passed through unchanged rotated_q, rotated_k = apply_pope_to_qk(pos_emb, q, k) assert rotated_q.shape == (2, 4, 256, 96) # 32*2 + 32 = 96 # Using ReLU instead of softplus as magnitude activation rotated_q2, rotated_k2 = apply_pope_to_qk(pos_emb, q, k, to_magnitude=F.relu) # Return complex tensors instead of interleaved real rotated_q_c, rotated_k_c = apply_pope_to_qk(pos_emb, q, k, return_complex=True) assert rotated_q_c.is_complex() ``` ``` -------------------------------- ### Compute Attention Similarity with PoPE Source: https://github.com/hunterhogan/pope-pytorch/blob/main/README.md Compute attention similarity using PoPE, avoiding expansion to avoid memory issues. ```python import torch from PoPE_pytorch import PoPE, compute_attn_similarity # define pope pope = PoPE(dim = 64, heads = 8).cuda() # get rotations pos_emb = pope(1024) # queries and keys q = torch.randn(1, 8, 1024, 64).cuda() k = torch.randn(1, 8, 1024, 64).cuda() # fused attention similarity, avoiding expanding 64 to 128 sim = compute_attn_similarity(q, k, pos_emb) # (1, 8, 1024, 1024) attn = sim.softmax(dim = -1) # the usual in attention.. ``` -------------------------------- ### PoPE Module Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt `PoPE` is a `torch.nn.Module` for computing polar-coordinate positional frequencies for 1D sequences. It can be instantiated and then called with a sequence length or position tensor to generate positional embeddings. It also provides a method to apply these embeddings to query and key tensors. ```APIDOC ## PoPE Module ### Description `PoPE` is a `torch.nn.Module` that computes polar-coordinate positional frequencies for 1D sequences. It stores learned inverse frequencies and a per-head learnable phase bias clamped to `[-2π, 0]`. Calling `forward` with a sequence length (or explicit position tensor) returns a `PolarEmbedReturn` named tuple of `(freqs, bias)` ready to be applied to queries and keys. ### Instantiate ```python import torch from PoPE_pytorch import PoPE # dim: per-head rotation dimension (must be <= head_dim in Q/K) # heads: number of attention heads pope = PoPE(dim=64, heads=8, theta=10000) ``` ### Generate Positional Embeddings ```python # Generate positional embeddings for a sequence of length 1024 pos_emb = pope(1024) # pos_emb.freqs shape: (1024, 64) # pos_emb.bias shape: (8, 64) ``` ### Apply to Q and K (Training: Full Context) ```python q = torch.randn(1, 8, 1024, 64) k = torch.randn(1, 8, 1024, 64) rotated_q, rotated_k = pope.apply_pope_to_qk(pos_emb, q, k) assert rotated_q.shape == (1, 8, 1024, 128) # dim doubles after polar interleaving assert rotated_k.shape == (1, 8, 1024, 128) ``` ### Apply to Q and K (Inference: Single New Token, Cached Keys) ```python rotated_q_inf, rotated_k_inf = pope.apply_pope_to_qk(pos_emb, q[..., -1:, :], k) assert rotated_q_inf.shape == (1, 8, 1, 128) assert rotated_k_inf.shape == (1, 8, 1024, 128) ``` ### Use Explicit Position Tensor ```python # Use explicit position tensor (e.g. for packed/offset sequences) positions = torch.arange(512, 512 + 1024).float() pos_emb_offset = pope(positions) ``` ``` -------------------------------- ### Compute Attention Similarity with PoPE Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt Computes the attention similarity matrix, optionally using a fused Triton kernel on CUDA for efficiency. Supports grouped query attention and different head layouts. ```python import torch from PoPE_pytorch import PoPE, compute_attn_similarity pope = PoPE(dim=64, heads=8).cuda() pos_emb = pope(1024) q = torch.randn(1, 8, 1024, 64).cuda() k = torch.randn(1, 8, 1024, 64).cuda() # Fused similarity (uses Triton kernel on CUDA) sim = compute_attn_similarity(q, k, pos_emb) # (1, 8, 1024, 1024) attn = sim.softmax(dim=-1) ``` ```python # Inference: single query token against full key cache q_single = torch.randn(1, 8, 1, 64).cuda() sim_single = compute_attn_similarity(q_single, k, pos_emb) # (1, 8, 1, 1024) ``` ```python # Grouped Query Attention (4 KV heads, 8 Q heads) pope_gqa = PoPE(dim=64, heads=4).cuda() pos_emb_gqa = pope_gqa(1024) q_gqa = torch.randn(1, 8, 1024, 64).cuda() k_gqa = torch.randn(1, 4, 1024, 64).cuda() sim_gqa = compute_attn_similarity(q_gqa, k_gqa, pos_emb_gqa) # (1, 8, 1024, 1024) ``` ```python # Layout: heads-last (batch, seq, heads, dim) q_hl = torch.randn(1, 1024, 8, 64).cuda() k_hl = torch.randn(1, 1024, 8, 64).cuda() sim_hl = compute_attn_similarity(q_hl, k_hl, pos_emb, head_dimension_at_first=False) ``` -------------------------------- ### compute_attn_similarity Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt Computes the attention similarity matrix after applying PoPE rotations. It can use a fused kernel on CUDA for efficiency or fall back to a standard einsum. Supports Grouped Query Attention (GQA). ```APIDOC ## compute_attn_similarity ### Description Computes the full `(batch, heads, seq_q, seq_k)` attention similarity matrix after applying PoPE rotations. On CUDA with Triton available it uses a fused kernel that avoids expanding head_dim to 2×head_dim; it falls back to a standard einsum otherwise. Supports grouped query attention (GQA) when `q.heads > k.heads`. ### Method `compute_attn_similarity(q, k, pos_emb, head_dimension_at_first=True)` ### Parameters - **q** (torch.Tensor) - Query tensor. - **k** (torch.Tensor) - Key tensor. - **pos_emb** (torch.Tensor) - Positional embedding tensor. - **head_dimension_at_first** (bool, optional) - Specifies if the head dimension is the first dimension after batch. Defaults to True. ### Request Example ```python import torch from PoPE_pytorch import PoPE, compute_attn_similarity pope = PoPE(dim=64, heads=8).cuda() pos_emb = pope(1024) q = torch.randn(1, 8, 1024, 64).cuda() k = torch.randn(1, 8, 1024, 64).cuda() # Fused similarity (uses Triton kernel on CUDA) sim = compute_attn_similarity(q, k, pos_emb) # (1, 8, 1024, 1024) attn = sim.softmax(dim=-1) # Inference: single query token against full key cache q_single = torch.randn(1, 8, 1, 64).cuda() sim_single = compute_attn_similarity(q_single, k, pos_emb) # (1, 8, 1, 1024) # Grouped Query Attention (4 KV heads, 8 Q heads) pope_gqa = PoPE(dim=64, heads=4).cuda() pos_emb_gqa = pope_gqa(1024) q_gqa = torch.randn(1, 8, 1024, 64).cuda() k_gqa = torch.randn(1, 4, 1024, 64).cuda() sim_gqa = compute_attn_similarity(q_gqa, k_gqa, pos_emb_gqa) # (1, 8, 1024, 1024) # Layout: heads-last (batch, seq, heads, dim) q_hl = torch.randn(1, 1024, 8, 64).cuda() k_hl = torch.randn(1, 1024, 8, 64).cuda() sim_hl = compute_attn_similarity(q_hl, k_hl, pos_emb, head_dimension_at_first=False) ``` ### Response #### Success Response (200) - **sim** (torch.Tensor) - The computed attention similarity matrix. ``` === COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.