Kernels:

kernels-community
/

flash-attn2

Trusted publisher

Kernel card Files Files and versions

xet

Community

kernels-bot commited on about 20 hours ago

Commit

571c42b

verified ·

1 Parent(s): 568865e

Uploaded using `kernel-builder`.

Browse files

Files changed (35) hide show

build/torch-stable-abi210-cu126-aarch64-linux/__init__.py +424 -0
build/torch-stable-abi210-cu126-aarch64-linux/_flash_attn2_cuda_460047c.abi3.so +3 -0
build/torch-stable-abi210-cu126-aarch64-linux/_ops.py +9 -0
build/torch-stable-abi210-cu126-aarch64-linux/flash_attn_interface.py +1708 -0
build/torch-stable-abi210-cu126-aarch64-linux/metadata.json +37 -0
build/torch-stable-abi210-cu126-aarch64-linux/metadata.json.sigstore +1 -0
build/torch-stable-abi210-cu126-aarch64-linux/symbols.json +1273 -0
build/torch-stable-abi210-cu128-aarch64-linux/__init__.py +424 -0
build/torch-stable-abi210-cu128-aarch64-linux/_flash_attn2_cuda_460047c.abi3.so +3 -0
build/torch-stable-abi210-cu128-aarch64-linux/_ops.py +9 -0
build/torch-stable-abi210-cu128-aarch64-linux/flash_attn_interface.py +1708 -0
build/torch-stable-abi210-cu128-aarch64-linux/metadata.json +39 -0
build/torch-stable-abi210-cu128-aarch64-linux/metadata.json.sigstore +1 -0
build/torch-stable-abi210-cu128-aarch64-linux/symbols.json +1273 -0
build/torch-stable-abi210-cu130-aarch64-linux/__init__.py +424 -0
build/torch-stable-abi210-cu130-aarch64-linux/_flash_attn2_cuda_460047c.abi3.so +3 -0
build/torch-stable-abi210-cu130-aarch64-linux/_ops.py +9 -0
build/torch-stable-abi210-cu130-aarch64-linux/flash_attn_interface.py +1708 -0
build/torch-stable-abi210-cu130-aarch64-linux/metadata.json +39 -0
build/torch-stable-abi210-cu130-aarch64-linux/metadata.json.sigstore +1 -0
build/torch-stable-abi210-cu130-aarch64-linux/symbols.json +1273 -0
build/torch-stable-abi210-cu132-aarch64-linux/__init__.py +424 -0
build/torch-stable-abi210-cu132-aarch64-linux/_flash_attn2_cuda_460047c.abi3.so +3 -0
build/torch-stable-abi210-cu132-aarch64-linux/_ops.py +9 -0
build/torch-stable-abi210-cu132-aarch64-linux/flash_attn_interface.py +1708 -0
build/torch-stable-abi210-cu132-aarch64-linux/metadata.json +39 -0
build/torch-stable-abi210-cu132-aarch64-linux/metadata.json.sigstore +1 -0
build/torch-stable-abi210-cu132-aarch64-linux/symbols.json +1273 -0
build/torch-stable-abi210-cu134-aarch64-linux/__init__.py +424 -0
build/torch-stable-abi210-cu134-aarch64-linux/_flash_attn2_cuda_460047c.abi3.so +3 -0
build/torch-stable-abi210-cu134-aarch64-linux/_ops.py +9 -0
build/torch-stable-abi210-cu134-aarch64-linux/flash_attn_interface.py +1708 -0
build/torch-stable-abi210-cu134-aarch64-linux/metadata.json +39 -0
build/torch-stable-abi210-cu134-aarch64-linux/metadata.json.sigstore +1 -0
build/torch-stable-abi210-cu134-aarch64-linux/symbols.json +1273 -0

build/torch-stable-abi210-cu126-aarch64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,424 @@

+from typing import List, Optional
+import torch
+from . import flash_attn_interface
+from ._ops import ops as flash_attn_ops
+from .flash_attn_interface import (
+    # helper to set up the Philox random number generator for dropout
+    _philox_generator,
+    # reexported for convenience to be used in the high-level API functions
+    flash_attn_func,
+    flash_attn_kvpacked_func,
+    flash_attn_qkvpacked_func,
+    flash_attn_varlen_func,
+    flash_attn_varlen_kvpacked_func,
+    flash_attn_varlen_qkvpacked_func,
+    flash_attn_with_kvcache,
+)
+__all__ = [
+    # High-level API (with autograd support)
+    "flash_attn_func",
+    "flash_attn_kvpacked_func",
+    "flash_attn_qkvpacked_func",
+    "flash_attn_varlen_func",
+    "flash_attn_varlen_kvpacked_func",
+    "flash_attn_varlen_qkvpacked_func",
+    "flash_attn_with_kvcache",
+    # Upstream interface module. Its `_flash_attn_forward`, `_flash_attn_backward`,
+    # `_flash_attn_varlen_forward` and `_flash_attn_varlen_backward` are supported
+    # low-level entry points with upstream's signatures: Torch custom ops without
+    # autograd that return the LSE and are traceable by torch.compile. The autograd
+    # Functions above are built on top of them.
+    "flash_attn_interface",
+    # Low-level ops (no autograd)
+    "fwd",
+    "varlen_fwd",
+    "bwd",
+    "varlen_bwd",
+    "fwd_kvcache",
+]
+def fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    return_softmax: bool = False,
+    gen: Optional[torch.Generator] = None,
+) -> List[torch.Tensor]:
+    """
+    Forward pass for multi-head attention.
+    Args:
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        out: Optional output tensor, same shape as q
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        return_softmax: Whether to return softmax weights
+        gen: Optional random number generator
+    Returns:
+        List of tensors: [output, softmax_lse, (softmax if return_softmax)]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    gen = _philox_generator(q, p_dropout, gen)
+    return flash_attn_ops.fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+def varlen_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    block_table: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    max_seqlen_q: int = 0,
+    max_seqlen_k: int = 0,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    zero_tensors: bool = False,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    return_softmax: bool = False,
+    gen: Optional[torch.Generator] = None,
+) -> List[torch.Tensor]:
+    """
+    Forward pass for multi-head attention with variable sequence lengths.
+    Args:
+        q: Query tensor of shape [total_q, num_heads, head_size]
+        k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]
+        cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]
+        out: Optional output tensor of shape [total_q, num_heads, head_size]
+        seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size]
+        leftpad_k: Optional left padding for keys of shape [batch_size]
+        block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        max_seqlen_q: Maximum sequence length for queries
+        max_seqlen_k: Maximum sequence length for keys
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        zero_tensors: Whether to zero tensors before computation
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        return_softmax: Whether to return softmax weights
+        gen: Optional random number generator
+    Returns:
+        List of tensors: [output, softmax_lse, (softmax if return_softmax)]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    gen = _philox_generator(q, p_dropout, gen)
+    return flash_attn_ops.varlen_fwd(
+        q,
+        k,
+        v,
+        out,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_k,
+        leftpad_k,
+        block_table,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        p_dropout,
+        softmax_scale,
+        zero_tensors,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+def bwd(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor] = None,
+    dk: Optional[torch.Tensor] = None,
+    dv: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    deterministic: bool = False,
+    gen: Optional[torch.Generator] = None,
+    rng_state: Optional[torch.Tensor] = None,
+) -> List[torch.Tensor]:
+    """
+    Backward pass for multi-head attention.
+    Args:
+        dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]
+        softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]
+        dq: Optional gradient tensor for queries, same shape as q
+        dk: Optional gradient tensor for keys, same shape as k
+        dv: Optional gradient tensor for values, same shape as v
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        deterministic: Whether to use deterministic algorithms
+        gen: Optional random number generator
+        rng_state: Optional RNG state from forward pass
+    Returns:
+        List of tensors: [dq, dk, dv]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    return flash_attn_ops.bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        gen,
+        rng_state,
+    )
+def varlen_bwd(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    dq: Optional[torch.Tensor] = None,
+    dk: Optional[torch.Tensor] = None,
+    dv: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    max_seqlen_q: int = 0,
+    max_seqlen_k: int = 0,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    zero_tensors: bool = False,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    deterministic: bool = False,
+    gen: Optional[torch.Generator] = None,
+    rng_state: Optional[torch.Tensor] = None,
+) -> List[torch.Tensor]:
+    """
+    Backward pass for multi-head attention with variable sequence lengths.
+    Args:
+        dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]
+        softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]
+        cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]
+        cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]
+        dq: Optional gradient tensor for queries, same shape as q
+        dk: Optional gradient tensor for keys, same shape as k
+        dv: Optional gradient tensor for values, same shape as v
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        max_seqlen_q: Maximum sequence length for queries
+        max_seqlen_k: Maximum sequence length for keys
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        zero_tensors: Whether to zero tensors before computation
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        deterministic: Whether to use deterministic algorithms
+        gen: Optional random number generator
+        rng_state: Optional RNG state from forward pass
+    Returns:
+        List of tensors: [dq, dk, dv]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    return flash_attn_ops.varlen_bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        p_dropout,
+        softmax_scale,
+        zero_tensors,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        gen,
+        rng_state,
+    )
+def fwd_kvcache(
+    q: torch.Tensor,
+    kcache: torch.Tensor,
+    vcache: torch.Tensor,
+    k: Optional[torch.Tensor] = None,
+    v: Optional[torch.Tensor] = None,
+    seqlens_k: Optional[torch.Tensor] = None,
+    rotary_cos: Optional[torch.Tensor] = None,
+    rotary_sin: Optional[torch.Tensor] = None,
+    cache_batch_idx: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    block_table: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    is_rotary_interleaved: bool = False,
+    num_splits: int = 1,
+) -> List[torch.Tensor]:
+    """
+    Forward pass for multi-head attention with KV cache.
+    Args:
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]
+        v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]
+        seqlens_k: Optional sequence lengths for keys of shape [batch_size]
+        rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2]
+        rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2]
+        cache_batch_idx: Optional indices to index into the KV cache
+        leftpad_k: Optional left padding for keys of shape [batch_size]
+        block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        out: Optional output tensor, same shape as q
+        softmax_scale: Scale factor for softmax
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        is_rotary_interleaved: Whether rotary embeddings are interleaved
+        num_splits: Number of splits for computation
+    Returns:
+        List of tensors: [output, softmax_lse]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    return flash_attn_ops.fwd_kvcache(
+        q,
+        kcache,
+        vcache,
+        k,
+        v,
+        seqlens_k,
+        rotary_cos,
+        rotary_sin,
+        cache_batch_idx,
+        leftpad_k,
+        block_table,
+        alibi_slopes,
+        out,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        is_rotary_interleaved,
+        num_splits,
+    )

build/torch-stable-abi210-cu126-aarch64-linux/_flash_attn2_cuda_460047c.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:de41e0887afd71fca22cc0b2f75a0fa132461ee98943529fb073f6767a33056b
+size 446708200

build/torch-stable-abi210-cu126-aarch64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _flash_attn2_cuda_460047c
+ops = torch.ops._flash_attn2_cuda_460047c
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn2_cuda_460047c::{op_name}"

build/torch-stable-abi210-cu126-aarch64-linux/flash_attn_interface.py ADDED Viewed

	@@ -0,0 +1,1708 @@

+# Copyright (c) 2023, Tri Dao.
+from typing import Optional, Sequence, Tuple, Union
+import torch
+import torch.nn as nn
+import os
+# # isort: off
+# # We need to import the CUDA kernels after importing torch
+# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE"
+# if USE_TRITON_ROCM:
+#     from .flash_attn_triton_amd import interface_fa as flash_attn
+# else:
+#     import flash_attn_2_cuda as flash_attn
+from ._ops import ops as flash_attn
+from ._ops import add_op_namespace_prefix
+# # isort: on
+# Torch >= 2.13 can pass a Generator through the stable ABI. Older versions reject the argument
+# before the kernel runs, so there the CUDA op has to draw its Philox seed and offset itself.
+_STABLE_ABI_GENERATOR = torch.__version__ >= "2.13"
+def maybe_contiguous(x):
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+def _get_device():
+    if torch.xpu.is_available():
+        return "xpu"
+    elif torch.cuda.is_available():
+        return "cuda"
+    else:
+        return "cpu"
+def _get_block_size_n(device, head_dim, is_dropout, is_causal):
+    assert head_dim <= 256
+    if device.type == "xpu":
+        if head_dim <= 96:
+            return 64
+        elif head_dim <= 128:
+            return 32
+        elif head_dim <= 256:
+            return 64
+        else:
+            return 32
+    # This should match the block sizes in the CUDA kernel
+    major, minor = torch.cuda.get_device_capability(device)
+    is_sm8x = (
+        major == 8 and minor > 0
+    )  # Only include sm86 and sm89, exclude sm80 (A100)
+    is_sm80 = major == 8 and minor == 0
+    is_sm90 = major == 9 and minor == 0
+    if head_dim <= 32:
+        return 128
+    if head_dim <= 64:
+        return 128 if not is_dropout else 64
+    elif head_dim <= 96:
+        return 64
+    elif head_dim <= 128:
+        if is_sm8x:
+            return 64 if (not is_dropout and is_causal) else 32
+        else:
+            return 64 if not is_dropout else 32
+    elif head_dim <= 192:
+        return 64
+    elif head_dim <= 224:
+        return 64
+    elif head_dim <= 256:
+        return 64
+def round_multiple(x, m):
+    return (x + m - 1) // m * m
+def _philox_generator(
+    q: torch.Tensor, dropout_p: float, generator: Optional[torch.Generator] = None
+) -> Optional[torch.Generator]:
+    """The generator to hand to `fwd` / `varlen_fwd`.
+    With dropout on CUDA this is the default generator unless one is given, so the kernel reserves
+    its seed and offset exactly like the non-stable build (see set_params_philox in flash_api.cpp).
+    """
+    if not q.is_cuda:
+        return generator
+    if dropout_p <= 0:
+        return None
+    if not _STABLE_ABI_GENERATOR:
+        if generator is not None:
+            raise ValueError(
+                "flash_attn2: passing a generator to the CUDA kernel requires torch >= 2.13; "
+                "seed the default generator with torch.manual_seed instead"
+            )
+        return None
+    if generator is None:
+        generator = torch.cuda.default_generators[q.device.index]
+    return generator
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_forward"),
+    mutates_args=(),
+    device_types=_get_device(),
+    # Tells torch.compile the forward draws random numbers, so it is never deduplicated or
+    # recomputed (e.g. under activation checkpointing).
+    tags=(torch.Tag.nondeterministic_seeded,),
+)
+def _flash_attn_forward(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    return_softmax: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    gen = _philox_generator(q, dropout_p)
+    out, softmax_lse, S_dmask, rng_state = flash_attn.fwd(
+        q,
+        k,
+        v,
+        None,
+        alibi_slopes,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out, softmax_lse, S_dmask, rng_state
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward"))
+def _flash_attn_forward_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    return_softmax: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    batch_size, seqlen_q, num_heads, head_size = q.shape
+    seqlen_k = k.shape[1]
+    out = torch.empty_like(q)
+    softmax_lse = torch.empty(
+        (batch_size, num_heads, seqlen_q),
+        dtype=torch.float32,
+        device=q.device,
+        layout=q.layout,
+    )
+    p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout)
+    if return_softmax:
+        p = torch.empty(
+            (
+                batch_size,
+                num_heads,
+                round_multiple(seqlen_q, 128),
+                round_multiple(seqlen_k, 128),
+            ),
+            dtype=q.dtype,
+            device=q.device,
+            layout=q.layout,
+        )
+    rng_state = torch.empty((2,), dtype=torch.int64, device=q.device)
+    return out, softmax_lse, p, rng_state
+_wrapped_flash_attn_forward = flash_attn._flash_attn_forward  # upstream alias, see CARD.md
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_varlen_forward"),
+    mutates_args=(),
+    device_types=_get_device(),
+    tags=(torch.Tag.nondeterministic_seeded,),
+)
+def _flash_attn_varlen_forward(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    return_softmax: bool = False,
+    block_table: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    gen = _philox_generator(q, dropout_p)
+    out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd(
+        q,
+        k,
+        v,
+        None,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_k,
+        leftpad_k,
+        block_table,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        zero_tensors,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    # if out.isnan().any() or softmax_lse.isnan().any():
+    #     breakpoint()
+    return out, softmax_lse, S_dmask, rng_state
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_forward"))
+def _flash_attn_varlen_forward_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    return_softmax: bool = False,
+    block_table: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    paged_kv = block_table is not None
+    batch_size = cu_seqlens_q.numel() - 1
+    total_q, num_heads, _ = q.shape
+    out = torch.empty_like(q)
+    softmax_lse = torch.empty(
+        (num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout
+    )
+    p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout)
+    seqlen_q_rounded = round_multiple(max_seqlen_q, 128)
+    seqlen_k_rounded = round_multiple(max_seqlen_k, 128)
+    if return_softmax:
+        p = torch.empty(
+            (batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded),
+            dtype=q.dtype,
+            device=q.device,
+            layout=q.layout,
+        )
+    rng_state = torch.empty((2,), dtype=torch.int64, device=q.device)
+    return out, softmax_lse, p, rng_state
+_wrapped_flash_attn_varlen_forward = flash_attn._flash_attn_varlen_forward  # upstream alias, see CARD.md
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_backward"),
+    mutates_args=("dq", "dk", "dv"),
+    device_types=_get_device(),
+)
+def _flash_attn_backward(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    # dq, dk, dv are allocated by us so they should already be contiguous
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    (
+        dq,
+        dk,
+        dv,
+        softmax_d,
+    ) = flash_attn.bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        alibi_slopes,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        None,
+        rng_state,
+    )
+    return softmax_d
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward"))
+def _flash_attn_backward_fake(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    if dq is None:
+        dq = torch.empty_like(q)
+    if dk is None:
+        dk = torch.empty_like(k)
+    if dv is None:
+        dv = torch.empty_like(v)
+    batch_size, seqlen_q, num_heads, _ = q.shape
+    softmax_d = torch.empty(
+        (batch_size, num_heads, round_multiple(seqlen_q, 128)),
+        device=q.device,
+        dtype=torch.float32,
+    )
+    return softmax_d
+_wrapped_flash_attn_backward = flash_attn._flash_attn_backward  # upstream alias, see CARD.md
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_varlen_backward"),
+    mutates_args=("dq", "dk", "dv"),
+    device_types=_get_device(),
+)
+def _flash_attn_varlen_backward(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> torch.Tensor:
+    # dq, dk, dv are allocated by us so they should already be contiguous
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    (
+        dq,
+        dk,
+        dv,
+        softmax_d,
+    ) = flash_attn.varlen_bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        zero_tensors,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        None,
+        rng_state,
+    )
+    # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any():
+    #     breakpoint()
+    return softmax_d
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_backward"))
+def _flash_attn_varlen_backward_fake(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> torch.Tensor:
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    batch_size = cu_seqlens_q.numel() - 1
+    total_q, num_heads, _ = q.shape
+    if dq is None:
+        dq = torch.empty_like(q)
+    if dk is None:
+        dk = torch.empty_like(k)
+    if dv is None:
+        dv = torch.empty_like(v)
+    softmax_d = torch.empty(
+        (num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32
+    )
+    return softmax_d
+_wrapped_flash_attn_varlen_backward = flash_attn._flash_attn_varlen_backward  # upstream alias, see CARD.md
+class FlashAttnQKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        qkv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and qkv.requires_grad
+        if softmax_scale is None:
+            softmax_scale = qkv.shape[-1] ** (-0.5)
+        q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach()
+        head_size_og = q.size(3)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+        )
+        if is_grad:
+            ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
+        qkv_shape = q.shape[:-2] + (3, *q.shape[-2:])
+        dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device)
+        head_size_og = dout.size(3)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dqkv[:, :, 0],
+            dqkv[:, :, 1],
+            dqkv[:, :, 2],
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dqkv = dqkv[..., : dout.shape[-1]]  # We could have padded the head dimension
+        return dqkv, None, None, None, None, None, None, None, None, None
+class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        qkv,
+        cu_seqlens,
+        max_seqlen,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and qkv.requires_grad
+        if softmax_scale is None:
+            softmax_scale = qkv.shape[-1] ** (-0.5)
+        q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach()
+        head_size_og = q.size(2)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
+            q,
+            k,
+            v,
+            cu_seqlens,
+            cu_seqlens,
+            max_seqlen,
+            max_seqlen,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+            block_table=None,
+        )
+        if is_grad:
+            ctx.save_for_backward(
+                q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state
+            )
+            ctx.dropout_p = dropout_p
+            ctx.max_seqlen = max_seqlen
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors
+        qkv_shape = q.shape[:-2] + (3, *q.shape[-2:])
+        dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device)
+        head_size_og = dout.size(2)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_varlen_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dqkv[:, 0],
+            dqkv[:, 1],
+            dqkv[:, 2],
+            cu_seqlens,
+            cu_seqlens,
+            ctx.max_seqlen,
+            ctx.max_seqlen,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dqkv = dqkv[..., : dout.shape[-1]]  # We could have padded the head dimension
+        return dqkv, None, None, None, None, None, None, None, None, None, None, None
+class FlashAttnKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        kv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach()
+        head_size_og = q.size(3)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+        )
+        if is_grad:
+            ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
+        dq = torch.empty_like(q)
+        kv_shape = k.shape[:-2] + (2, *k.shape[-2:])
+        dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device)
+        head_size_og = dout.size(3)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dkv[:, :, 0],
+            dkv[:, :, 1],
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dkv = dkv[..., : dout.shape[-1]]
+        return dq, dkv, None, None, None, None, None, None, None, None, None
+class FlashAttnVarlenKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        kv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        k, v = kv[:, 0].detach(), kv[:, 1].detach()
+        head_size_og = q.size(2)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+            block_table=None,
+        )
+        if is_grad:
+            ctx.save_for_backward(
+                q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state
+            )
+            ctx.dropout_p = dropout_p
+            ctx.max_seqlen_q = max_seqlen_q
+            ctx.max_seqlen_k = max_seqlen_k
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = (
+            ctx.saved_tensors
+        )
+        dq = torch.empty_like(q)
+        kv_shape = k.shape[:-2] + (2, *k.shape[-2:])
+        dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device)
+        head_size_og = dout.size(2)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_varlen_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dkv[:, 0],
+            dkv[:, 1],
+            cu_seqlens_q,
+            cu_seqlens_k,
+            ctx.max_seqlen_q,
+            ctx.max_seqlen_k,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dkv = dkv[..., : dout.shape[-1]]
+        return (
+            dq,
+            dkv,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+class FlashAttnFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        head_size_og = q.size(3)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+        )
+        if is_grad:
+            ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
+        dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+        head_size_og = dout.size(3)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dk,
+            dv,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dk = dk[..., : dout.shape[-1]]
+        dv = dv[..., : dout.shape[-1]]
+        return dq, dk, dv, None, None, None, None, None, None, None, None, None
+class FlashAttnVarlenFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        block_table,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        head_size_og = q.size(2)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+            block_table=block_table,
+        )
+        if is_grad:
+            ctx.save_for_backward(
+                q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state
+            )
+            ctx.dropout_p = dropout_p
+            ctx.max_seqlen_q = max_seqlen_q
+            ctx.max_seqlen_k = max_seqlen_k
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = (
+            ctx.saved_tensors
+        )
+        dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+        head_size_og = dout.size(2)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_varlen_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dk,
+            dv,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            ctx.max_seqlen_q,
+            ctx.max_seqlen_k,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dk = dk[..., : dout.shape[-1]]
+        dv = dv[..., : dout.shape[-1]]
+        return (
+            dq,
+            dk,
+            dv,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+def flash_attn_qkvpacked_func(
+    qkv,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # <=0.0 means deactivate
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If Q, K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of Q, K, V.
+    For multi-query and grouped-query attention (MQA/GQA), please see
+    flash_attn_kvpacked_func and flash_attn_func.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.
+    Arguments:
+        qkv: (batch_size, seqlen, 3, nheads, headdim)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to
+            the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnQKVPackedFunc.apply(
+        qkv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_kvpacked_func(
+    q,
+    kv,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of K, V.
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        kv: (batch_size, seqlen, 2, nheads_k, headdim)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnKVPackedFunc.apply(
+        q,
+        kv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_func(
+    q,
+    k,
+    v,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        k: (batch_size, seqlen, nheads_k, headdim)
+        v: (batch_size, seqlen, nheads_k, headdim)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnFunc.apply(
+        q,
+        k,
+        v,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_varlen_qkvpacked_func(
+    qkv,
+    cu_seqlens,
+    max_seqlen,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If Q, K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_varlen_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of Q, K, V.
+    For multi-query and grouped-query attention (MQA/GQA), please see
+    flash_attn_varlen_kvpacked_func and flash_attn_varlen_func.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.
+    Arguments:
+        qkv: (total, 3, nheads, headdim), where total = total number of tokens in the batch.
+        cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into qkv.
+        max_seqlen: int. Maximum sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnVarlenQKVPackedFunc.apply(
+        qkv,
+        cu_seqlens,
+        max_seqlen,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_varlen_kvpacked_func(
+    q,
+    kv,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of K, V.
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.
+        kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into q.
+        cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into kv.
+        max_seqlen_q: int. Maximum query sequence length in the batch.
+        max_seqlen_k: int. Maximum key sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnVarlenKVPackedFunc.apply(
+        q,
+        kv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_varlen_func(
+    q,
+    k,
+    v,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+    block_table=None,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.
+        k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into q.
+        cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into kv.
+        max_seqlen_q: int. Maximum query sequence length in the batch.
+        max_seqlen_k: int. Maximum key sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnVarlenFunc.apply(
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        block_table,
+        False if q.device.type == "cpu" else torch.is_grad_enabled(),
+    )
+def flash_attn_with_kvcache(
+    q,
+    k_cache,
+    v_cache,
+    k=None,
+    v=None,
+    rotary_cos=None,
+    rotary_sin=None,
+    cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None,
+    cache_batch_idx: Optional[torch.Tensor] = None,
+    cache_leftpad: Optional[torch.Tensor] = None,
+    block_table: Optional[torch.Tensor] = None,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    rotary_interleaved=True,
+    alibi_slopes=None,
+    num_splits=0,
+    return_softmax_lse=False,
+):
+    """
+    If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from
+    k and v. This is useful for incremental decoding: you can pass in the cached keys/values from
+    the previous step, and update them with the new keys/values from the current step, and do
+    attention with the updated cache, all in 1 kernel.
+    If you pass in k / v, you must make sure that the cache is large enough to hold the new values.
+    For example, the KV cache could be pre-allocated with the max sequence length, and you can use
+    cache_seqlens to keep track of the current sequence lengths of each sequence in the batch.
+    Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be
+    rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
+    If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos
+    and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
+    If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at
+    indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens).
+    See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function.
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Note: Does not support backward pass.
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,
+            or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)
+            page_block_size must be a multiple of 256.
+        v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,
+            or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)
+        k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate
+            k with k_cache, starting at the indices specified by cache_seqlens.
+        v [optional]: (batch_size, seqlen_new, nheads_k, headdim). Similar to k.
+        rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding
+            to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16.
+        rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos.
+        cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the
+            KV cache.
+        cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache.
+            If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1].
+            If the indices are not distinct, and k and v are provided, the values updated in the cache
+                 might come from any of the duplicate indices.
+        cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0.
+        block_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in.
+            If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False,
+            rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1
+            (i.e. GPT-NeoX style).
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        num_splits: int. If > 1, split the key/value into this many chunks along the sequence.
+           If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic
+           to automatically determine the number of splits.
+           Don't change this unless you know what you are doing.
+        return_softmax_lse: bool. Whether to return the logsumexp of the attention scores.
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+    """
+    assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension"
+    assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension"
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** (-0.5)
+    if cache_seqlens is not None and isinstance(cache_seqlens, int):
+        cache_seqlens = torch.full(
+            (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device
+        )
+        cache_seqlens = maybe_contiguous(cache_seqlens)
+    cache_batch_idx = maybe_contiguous(cache_batch_idx)
+    block_table = maybe_contiguous(block_table)
+    out, softmax_lse = flash_attn.fwd_kvcache(
+        q,
+        k_cache,
+        v_cache,
+        k,
+        v,
+        cache_seqlens,
+        rotary_cos,
+        rotary_sin,
+        cache_batch_idx,
+        cache_leftpad,
+        block_table,
+        alibi_slopes,
+        None,
+        softmax_scale,
+        causal,
+        window_size[0],
+        window_size[1],
+        softcap,
+        rotary_interleaved,
+        num_splits,
+    )
+    return (out, softmax_lse) if return_softmax_lse else out

build/torch-stable-abi210-cu126-aarch64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,37 @@

+{
+  "name": "flash-attn2",
+  "id": "_flash_attn2_cuda_460047c",
+  "version": 4,
+  "kernels-minver": "0.14.0",
+  "license": "BSD-3-Clause",
+  "python-depends": [],
+  "kernel-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.0",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Nz8hqrusjoUHQqszpxfF2USMwUXn1VfpvcmtAAvutSM=",
+      "_flash_attn2_cuda_460047c.abi3.so": "3kHgiHr9cfyiLMCy91oPoTJGHumJQ1KfsHP2dnozBWs=",
+      "_ops.py": "DV4Eo4+s3aaNmDsfXoZiYKNLnRUJqlsdDGlMsdKFgk4=",
+      "flash_attn_interface.py": "6UYUhfUax4+yOUygyS7K8dVjrIjU5GQr1AJ2gSp8rPo=",
+      "symbols.json": "uvC+87tHUhtHv7L/Nl53T/Ms7ceX317/jwQBoqHY14A="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.18.0-dev0",
+      "commit": "d054906a13710b811b96697b4635f03020842733",
+      "dirty": false
+    },
+    "kernel": {
+      "commit": "460047c291fc7f0b3dbd4239b46bd5e1281b1806",
+      "dirty": false
+    }
+  }
+}

build/torch-stable-abi210-cu126-aarch64-linux/metadata.json.sigstore ADDED Viewed

	@@ -0,0 +1 @@

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrGgAwIBAgIUBo1ks8d30sjWI29uW4tvs06Uni0wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDI0MDAwNzQ2WhcNMjYxMDI0MDAxNzQ2WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEEzA99VPpsYDB2mkOR8D8Ny5mxwdRyzlpQV+P+bE8CoxdCpWd09IH013MGaePr13ZWqTY7DEQ8+3bJoMtXYny1aOCBdAwggXMMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBCg0NjAwNDdjMjkxZmM3ZjBiM2RiZDQyMzliNDZiZDVlMTI4MWIxODA2MBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDCg0NjAwNDdjMjkxZmM3ZjBiM2RiZDQyMzliNDZiZDVlMTI4MWIxODA2MBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDCg0NjAwNDdjMjkxZmM3ZjBiM2RiZDQyMzliNDZiZDVlMTI4MWIxODA2MB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoNDYwMDQ3YzI5MWZjN2YwYjNkYmQ0MjM5YjQ2YmQ1ZTEyODFiMTgwNjAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTk4MjY0MDUxNS9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGJBgorBgEEAdZ5AgQCBHsEeQB3AHUA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhazwlcQAABAMARjBEAiAUZQB6BhKRvHss/V7Ey1BQO2kpXcz9jH8m+ZAP1V5LFAIgTsU5KmrI4NC2ppUGMdlYZpj4qquMR2SyJjGHNO4xscMwCgYIKoZIzj0EAwMDaQAwZgIxAJc9aIUeNt/gf8rW27S6DBEFtz6OSP1nRPBySOWm/tr6P9XqqJbo1SB24NH0Ktrx+gIxAMPeNst7UVdTnMnp1aNXEoSk+PGQfGkM6IwI7YgGGlhPY/PKjVFPEgFfZvpFBP4P9Q=="},"tlogEntries":[{"logIndex":"3430230245","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792800466","inclusionPromise":{"signedEntryTimestamp":"MEQCICEYIjYMuYPLk7awtJYUASuf8XztXLMTPs4X+7Lpm3SDAiBQ2efj7gnDF/2NSMrJ86bJTfQ7rYpmYzchaNwDe1o2Jg=="},"inclusionProof":{"logIndex":"3308325983","rootHash":"9RhrJeVdTN1XWyq96fL3jjoOsc1nFsM/HES3FIUszPY=","treeSize":"3308326223","hashes":["ZZUqFo+YHPeC5R2p2SNck7J01R9ENmB4BWVMlc/slXA=","w1p4nI+R9aogwUiO35DVWRpok9RtvnHgMp91m1i7zoQ=","+kPgoLboWgPIQiwgl1SZtK3PTOu2+ACMpj/nlyywJEQ=","Asr5xdoDBoe5g4/E1xBri09H+TR7Hg8RYo8YPR50xUo=","s/j4WllkZVu4yLXws1v72aG8K1grAGGvb8ScmOZRxos=","JJNA2oVUYuX1sAXI7HnJAGnj/FjXacYPyZRxdAttFwg=","TjxCovIPh2dORfx6kgB6177znXUUsc7rfrk5J/CY1Ig=","UJUvB1WAmOcG2neywwATlbUwjcr87A3PcVOC1+jgJC4=","pUGUimNXW8f8DrPF31Qf+xxBakJshScZqbbDtWrHgkA=","PuwdRME/93l6s2BMRjCojmjsVC29KnBqdtPSPE1l6cQ=","tWdxtMrYqmZ0zKHnOCAERPP7SvuDvubtw3jwj7rH4uA=","8x0N/oAmvpY+6DQUGTU+i6a1Cuqphro2WbV5WZ/RHL8=","DMchGrSAeKV9qHo75XJV3UFZ/CHF8zofPRu+NgVRwA4=","V6dXjeDcwYxBG/9gVSBM8rYdPrAqGpZsf1e7v6l67GA=","iDfImp65w8dirRfYGNgFA34JTUYDWbcJwNBBU/Ok+X4=","YYCr61A8azar3dfO8Ct+pri7LqDxkUlUkfs7fiqYErU=","1DaDKzrCNNBhmSdY5vFjz+NpAmxmYfYEgBCdPLrMP1A=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3308326223\n9RhrJeVdTN1XWyq96fL3jjoOsc1nFsM/HES3FIUszPY=\n\n— rekor.sigstore.dev wNI9ajBFAiEAwzn/01eNUZJXH1n17uP842XQLlrV1cE3GdpyvYbwm7wCICUaMVfm6QDW2KxHRr4EYs2/flTQ5yF/PekdNUp89Ots\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI0YWZmMjRjYWUwMzRiNWVlNWNkZGI2ZDg5MTI5ODQ5ZWEwZmI1ZWYyOWI5ZjZkMjMzN2FlZTJjZTc5NzRmNzY3In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJQmo4dXJpMERPaEdVM3pGVGJOeGRJU09HTVhDR3lCTXdBbkJtZGhzY0hrWUFpQVpXbWVnbFVSSGtXMm82WlRYUFVSajkvRWcrZ25pV0lDUmdEWTVsV0VGREE9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KSFowRjNTVUpCWjBsVlFtOHhhM000WkRNd2MycFhTVEk1ZFZjMGRIWnpNRFpWYm1rd2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFU1RCTlJFRjNUbnBSTWxkb1kwNU5hbGw0VFVSSk1FMUVRWGhPZWxFeVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZGZWtFNU9WWlFjSE5aUkVJeWJXdFBVamhFT0U1NU5XMTRkMlJTZVhwc2NGRldLMUFLSzJKRk9FTnZlR1JEY0Zka01EbEpTREF4TTAxSFlXVlFjakV6V2xkeFZGazNSRVZST0NzellrcHZUWFJZV1c1NU1XRlBRMEprUVhkbloxaE5UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmN3VG1wQmR3cE9SR1JxVFdwcmVGcHRUVE5hYWtKcFRUSlNhVnBFVVhsTmVteHBUa1JhYVZwRVZteE5WRWswVFZkSmVFOUVRVEpOUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEWnpCT2FrRjNUa1JrYWsxcWEzaGFiVTB6Q2xwcVFtbE5NbEpwV2tSUmVVMTZiR2xPUkZwcFdrUldiRTFVU1RSTlYwbDRUMFJCTWsxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEWnpCT2FrRjNDazVFWkdwTmFtdDRXbTFOTTFwcVFtbE5NbEpwV2tSUmVVMTZiR2xPUkZwcFdrUldiRTFVU1RSTlYwbDRUMFJCTWsxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlPUkZsM1RVUlJNMWw2U1RWTlYxcHFUakpaZDFscVRtdFpiVkV3VFdwTk5WbHFVVEpaYlZFeFdsUkZlVTlFUm1rS1RWUm5kMDVxUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUnJORTFxV1RCTlJGVjRUbE01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkS0NrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NITkZaVkZDTTBGSVZVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvWVhwM2JHTlJRVUZDUVUxQlVtcENSVUZwUVZWYVVVSTJRbWhMVW5aSWMzTXZWamRGZVRGQ1VVOHlhM0JZWTNvNWFrZzRiUW9yV2tGUU1WWTFURVpCU1dkVWMxVTFTMjF5U1RST1F6SndjRlZIVFdSc1dWcHdhalJ4Y1hWTlVqSlRlVXBxUjBoT1R6UjRjMk5OZDBObldVbExiMXBKQ25wcU1FVkJkMDFFWVZGQmQxcG5TWGhCU21NNVlVbFZaVTUwTDJkbU9ISlhNamRUTmtSQ1JVWjBlalpQVTFBeGJsSlFRbmxUVDFkdEwzUnlObEE1V0hFS2NVcGliekZUUWpJMFRrZ3dTM1J5ZUN0blNYaEJUVkJsVG5OME4xVldaRlJ1VFc1d01XRk9XRVZ2VTJzclVFZFJaa2RyVFRaSmQwazNXV2RIUjJ4b1VBcFpMMUJMYWxaR1VFVm5SbVphZG5CR1FsQTBVRGxSUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg4XFkww1s5reoRTvhYd+uS/V0ItJYWN0bJLqR8Yxc4awCFQDBdz2E5aFNDadtTqGVkrBUsXTEKxgPMjAyNjEwMjQwMDA3NDZaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAyNDAwMDc0NlowLwYJKoZIhvcNAQkEMSIEIE5cpwDezUJr7nCJtvtbmyvIluFy294vwygMYx+wZhO6MIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMB2jAvel7SgzAqKiewQFzpqBT+qoq78RFvWuJcDJyDg+G/XV9PI/W/T3+ZcfW0oUMAIxAPA+84KPzGjn//rq7+fWSuSF8WPE7JIj6PN4Ki5BnCwuCZiqK/6w8oszWjxiz+8tDA=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"Sv8kyuA0te5c3bbYkSmEnqD7XvKbn20jN67iznl092c="},"signature":"MEQCIBj8uri0DOhGU3zFTbNxdISOGMXCGyBMwAnBmdhscHkYAiAZWmeglURHkW2o6ZTXPURj9/Eg+gniWICRgDY5lWEFDA=="}}

build/torch-stable-abi210-cu126-aarch64-linux/symbols.json ADDED Viewed

	@@ -0,0 +1,1273 @@

+{
+  "schema_version": 1,
+  "module": "_flash_attn2_cuda_460047c",
+  "functions": [
+    {
+      "name": "flash_attn_func",
+      "qualname": "flash_attn_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (batch_size, seqlen, nheads, headdim)\n    k: (batch_size, seqlen, nheads_k, headdim)\n    v: (batch_size, seqlen, nheads_k, headdim)\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_kvpacked_func",
+      "qualname": "flash_attn_kvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of K, V.\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (batch_size, seqlen, nheads, headdim)\n    kv: (batch_size, seqlen, 2, nheads_k, headdim)\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "kv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_qkvpacked_func",
+      "qualname": "flash_attn_qkvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf Q, K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of Q, K, V.\nFor multi-query and grouped-query attention (MQA/GQA), please see\nflash_attn_kvpacked_func and flash_attn_func.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.\n\nArguments:\n    qkv: (batch_size, seqlen, 3, nheads, headdim)\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to\n        the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "qkv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_varlen_func",
+      "qualname": "flash_attn_varlen_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.\n    k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.\n    v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.\n    cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into q.\n    cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into kv.\n    max_seqlen_q: int. Maximum query sequence length in the batch.\n    max_seqlen_k: int. Maximum key sequence length in the batch.\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (total, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_varlen_kvpacked_func",
+      "qualname": "flash_attn_varlen_kvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of K, V.\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.\n    kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch.\n    cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into q.\n    cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into kv.\n    max_seqlen_q: int. Maximum query sequence length in the batch.\n    max_seqlen_k: int. Maximum key sequence length in the batch.\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (total, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "kv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_varlen_qkvpacked_func",
+      "qualname": "flash_attn_varlen_qkvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf Q, K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_varlen_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of Q, K, V.\nFor multi-query and grouped-query attention (MQA/GQA), please see\nflash_attn_varlen_kvpacked_func and flash_attn_varlen_func.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.\n\nArguments:\n    qkv: (total, 3, nheads, headdim), where total = total number of tokens in the batch.\n    cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into qkv.\n    max_seqlen: int. Maximum sequence length in the batch.\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (total, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "qkv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_with_kvcache",
+      "qualname": "flash_attn_with_kvcache",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from\nk and v. This is useful for incremental decoding: you can pass in the cached keys/values from\nthe previous step, and update them with the new keys/values from the current step, and do\nattention with the updated cache, all in 1 kernel.\n\nIf you pass in k / v, you must make sure that the cache is large enough to hold the new values.\nFor example, the KV cache could be pre-allocated with the max sequence length, and you can use\ncache_seqlens to keep track of the current sequence lengths of each sequence in the batch.\n\nAlso apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be\nrotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.\nIf causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos\nand rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.\nIf not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at\nindices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens).\n\nSee tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function.\n\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nNote: Does not support backward pass.\n\nArguments:\n    q: (batch_size, seqlen, nheads, headdim)\n    k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,\n        or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)\n        page_block_size must be a multiple of 256.\n    v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,\n        or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)\n    k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate\n        k with k_cache, starting at the indices specified by cache_seqlens.\n    v [optional]: (batch_size, seqlen_new, nheads_k, headdim). Similar to k.\n    rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding\n        to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16.\n    rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos.\n    cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the\n        KV cache.\n    cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache.\n        If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1].\n        If the indices are not distinct, and k and v are provided, the values updated in the cache\n             might come from any of the duplicate indices.\n    cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0.\n    block_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in.\n        If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False,\n        rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1\n        (i.e. GPT-NeoX style).\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    num_splits: int. If > 1, split the key/value into this many chunks along the sequence.\n       If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic\n       to automatically determine the number of splits.\n       Don't change this unless you know what you are doing.\n    return_softmax_lse: bool. Whether to return the logsumexp of the attention scores.\n\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k_cache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "v_cache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "rotary_cos",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "rotary_sin",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "cache_seqlens",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "int | torch.Tensor | None"
+          },
+          {
+            "name": "cache_batch_idx",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "cache_leftpad",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "rotary_interleaved",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "True",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "num_splits",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": null
+          },
+          {
+            "name": "return_softmax_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "fwd",
+      "qualname": "fwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Forward pass for multi-head attention.\n\nArgs:\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    out: Optional output tensor, same shape as q\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    return_softmax: Whether to return softmax weights\n    gen: Optional random number generator\n\nReturns:\n    List of tensors: [output, softmax_lse, (softmax if return_softmax)]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "return_softmax",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "varlen_fwd",
+      "qualname": "varlen_fwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Forward pass for multi-head attention with variable sequence lengths.\n\nArgs:\n    q: Query tensor of shape [total_q, num_heads, head_size]\n    k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]\n    cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]\n    out: Optional output tensor of shape [total_q, num_heads, head_size]\n    seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size]\n    leftpad_k: Optional left padding for keys of shape [batch_size]\n    block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    max_seqlen_q: Maximum sequence length for queries\n    max_seqlen_k: Maximum sequence length for keys\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    zero_tensors: Whether to zero tensors before computation\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    return_softmax: Whether to return softmax weights\n    gen: Optional random number generator\n\nReturns:\n    List of tensors: [output, softmax_lse, (softmax if return_softmax)]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "seqused_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "leftpad_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "zero_tensors",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "return_softmax",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "bwd",
+      "qualname": "bwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Backward pass for multi-head attention.\n\nArgs:\n    dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]\n    softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]\n    dq: Optional gradient tensor for queries, same shape as q\n    dk: Optional gradient tensor for keys, same shape as k\n    dv: Optional gradient tensor for values, same shape as v\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    deterministic: Whether to use deterministic algorithms\n    gen: Optional random number generator\n    rng_state: Optional RNG state from forward pass\n\nReturns:\n    List of tensors: [dq, dk, dv]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "dout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "softmax_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "dq",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dk",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          },
+          {
+            "name": "rng_state",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "varlen_bwd",
+      "qualname": "varlen_bwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Backward pass for multi-head attention with variable sequence lengths.\n\nArgs:\n    dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]\n    softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]\n    cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]\n    cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]\n    dq: Optional gradient tensor for queries, same shape as q\n    dk: Optional gradient tensor for keys, same shape as k\n    dv: Optional gradient tensor for values, same shape as v\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    max_seqlen_q: Maximum sequence length for queries\n    max_seqlen_k: Maximum sequence length for keys\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    zero_tensors: Whether to zero tensors before computation\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    deterministic: Whether to use deterministic algorithms\n    gen: Optional random number generator\n    rng_state: Optional RNG state from forward pass\n\nReturns:\n    List of tensors: [dq, dk, dv]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "dout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "softmax_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "dq",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dk",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "zero_tensors",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          },
+          {
+            "name": "rng_state",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "fwd_kvcache",
+      "qualname": "fwd_kvcache",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Forward pass for multi-head attention with KV cache.\n\nArgs:\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]\n    v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]\n    seqlens_k: Optional sequence lengths for keys of shape [batch_size]\n    rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2]\n    rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2]\n    cache_batch_idx: Optional indices to index into the KV cache\n    leftpad_k: Optional left padding for keys of shape [batch_size]\n    block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    out: Optional output tensor, same shape as q\n    softmax_scale: Scale factor for softmax\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    is_rotary_interleaved: Whether rotary embeddings are interleaved\n    num_splits: Number of splits for computation\n\nReturns:\n    List of tensors: [output, softmax_lse]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "kcache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "vcache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "rotary_cos",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "rotary_sin",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "cache_batch_idx",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "leftpad_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "is_rotary_interleaved",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "num_splits",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "1",
+            "annotation": "int"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    }
+  ],
+  "layers": []
+}

build/torch-stable-abi210-cu128-aarch64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,424 @@

+from typing import List, Optional
+import torch
+from . import flash_attn_interface
+from ._ops import ops as flash_attn_ops
+from .flash_attn_interface import (
+    # helper to set up the Philox random number generator for dropout
+    _philox_generator,
+    # reexported for convenience to be used in the high-level API functions
+    flash_attn_func,
+    flash_attn_kvpacked_func,
+    flash_attn_qkvpacked_func,
+    flash_attn_varlen_func,
+    flash_attn_varlen_kvpacked_func,
+    flash_attn_varlen_qkvpacked_func,
+    flash_attn_with_kvcache,
+)
+__all__ = [
+    # High-level API (with autograd support)
+    "flash_attn_func",
+    "flash_attn_kvpacked_func",
+    "flash_attn_qkvpacked_func",
+    "flash_attn_varlen_func",
+    "flash_attn_varlen_kvpacked_func",
+    "flash_attn_varlen_qkvpacked_func",
+    "flash_attn_with_kvcache",
+    # Upstream interface module. Its `_flash_attn_forward`, `_flash_attn_backward`,
+    # `_flash_attn_varlen_forward` and `_flash_attn_varlen_backward` are supported
+    # low-level entry points with upstream's signatures: Torch custom ops without
+    # autograd that return the LSE and are traceable by torch.compile. The autograd
+    # Functions above are built on top of them.
+    "flash_attn_interface",
+    # Low-level ops (no autograd)
+    "fwd",
+    "varlen_fwd",
+    "bwd",
+    "varlen_bwd",
+    "fwd_kvcache",
+]
+def fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    return_softmax: bool = False,
+    gen: Optional[torch.Generator] = None,
+) -> List[torch.Tensor]:
+    """
+    Forward pass for multi-head attention.
+    Args:
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        out: Optional output tensor, same shape as q
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        return_softmax: Whether to return softmax weights
+        gen: Optional random number generator
+    Returns:
+        List of tensors: [output, softmax_lse, (softmax if return_softmax)]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    gen = _philox_generator(q, p_dropout, gen)
+    return flash_attn_ops.fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+def varlen_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    block_table: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    max_seqlen_q: int = 0,
+    max_seqlen_k: int = 0,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    zero_tensors: bool = False,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    return_softmax: bool = False,
+    gen: Optional[torch.Generator] = None,
+) -> List[torch.Tensor]:
+    """
+    Forward pass for multi-head attention with variable sequence lengths.
+    Args:
+        q: Query tensor of shape [total_q, num_heads, head_size]
+        k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]
+        cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]
+        out: Optional output tensor of shape [total_q, num_heads, head_size]
+        seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size]
+        leftpad_k: Optional left padding for keys of shape [batch_size]
+        block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        max_seqlen_q: Maximum sequence length for queries
+        max_seqlen_k: Maximum sequence length for keys
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        zero_tensors: Whether to zero tensors before computation
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        return_softmax: Whether to return softmax weights
+        gen: Optional random number generator
+    Returns:
+        List of tensors: [output, softmax_lse, (softmax if return_softmax)]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    gen = _philox_generator(q, p_dropout, gen)
+    return flash_attn_ops.varlen_fwd(
+        q,
+        k,
+        v,
+        out,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_k,
+        leftpad_k,
+        block_table,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        p_dropout,
+        softmax_scale,
+        zero_tensors,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+def bwd(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor] = None,
+    dk: Optional[torch.Tensor] = None,
+    dv: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    deterministic: bool = False,
+    gen: Optional[torch.Generator] = None,
+    rng_state: Optional[torch.Tensor] = None,
+) -> List[torch.Tensor]:
+    """
+    Backward pass for multi-head attention.
+    Args:
+        dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]
+        softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]
+        dq: Optional gradient tensor for queries, same shape as q
+        dk: Optional gradient tensor for keys, same shape as k
+        dv: Optional gradient tensor for values, same shape as v
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        deterministic: Whether to use deterministic algorithms
+        gen: Optional random number generator
+        rng_state: Optional RNG state from forward pass
+    Returns:
+        List of tensors: [dq, dk, dv]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    return flash_attn_ops.bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        gen,
+        rng_state,
+    )
+def varlen_bwd(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    dq: Optional[torch.Tensor] = None,
+    dk: Optional[torch.Tensor] = None,
+    dv: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    max_seqlen_q: int = 0,
+    max_seqlen_k: int = 0,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    zero_tensors: bool = False,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    deterministic: bool = False,
+    gen: Optional[torch.Generator] = None,
+    rng_state: Optional[torch.Tensor] = None,
+) -> List[torch.Tensor]:
+    """
+    Backward pass for multi-head attention with variable sequence lengths.
+    Args:
+        dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]
+        softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]
+        cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]
+        cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]
+        dq: Optional gradient tensor for queries, same shape as q
+        dk: Optional gradient tensor for keys, same shape as k
+        dv: Optional gradient tensor for values, same shape as v
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        max_seqlen_q: Maximum sequence length for queries
+        max_seqlen_k: Maximum sequence length for keys
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        zero_tensors: Whether to zero tensors before computation
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        deterministic: Whether to use deterministic algorithms
+        gen: Optional random number generator
+        rng_state: Optional RNG state from forward pass
+    Returns:
+        List of tensors: [dq, dk, dv]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    return flash_attn_ops.varlen_bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        p_dropout,
+        softmax_scale,
+        zero_tensors,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        gen,
+        rng_state,
+    )
+def fwd_kvcache(
+    q: torch.Tensor,
+    kcache: torch.Tensor,
+    vcache: torch.Tensor,
+    k: Optional[torch.Tensor] = None,
+    v: Optional[torch.Tensor] = None,
+    seqlens_k: Optional[torch.Tensor] = None,
+    rotary_cos: Optional[torch.Tensor] = None,
+    rotary_sin: Optional[torch.Tensor] = None,
+    cache_batch_idx: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    block_table: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    is_rotary_interleaved: bool = False,
+    num_splits: int = 1,
+) -> List[torch.Tensor]:
+    """
+    Forward pass for multi-head attention with KV cache.
+    Args:
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]
+        v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]
+        seqlens_k: Optional sequence lengths for keys of shape [batch_size]
+        rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2]
+        rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2]
+        cache_batch_idx: Optional indices to index into the KV cache
+        leftpad_k: Optional left padding for keys of shape [batch_size]
+        block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        out: Optional output tensor, same shape as q
+        softmax_scale: Scale factor for softmax
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        is_rotary_interleaved: Whether rotary embeddings are interleaved
+        num_splits: Number of splits for computation
+    Returns:
+        List of tensors: [output, softmax_lse]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    return flash_attn_ops.fwd_kvcache(
+        q,
+        kcache,
+        vcache,
+        k,
+        v,
+        seqlens_k,
+        rotary_cos,
+        rotary_sin,
+        cache_batch_idx,
+        leftpad_k,
+        block_table,
+        alibi_slopes,
+        out,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        is_rotary_interleaved,
+        num_splits,
+    )

build/torch-stable-abi210-cu128-aarch64-linux/_flash_attn2_cuda_460047c.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:cb60d33341a0f1262fa11d919c0776398651e327cdfeb7d21fad3c2aeaf8513b
+size 1036209976

build/torch-stable-abi210-cu128-aarch64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _flash_attn2_cuda_460047c
+ops = torch.ops._flash_attn2_cuda_460047c
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn2_cuda_460047c::{op_name}"

build/torch-stable-abi210-cu128-aarch64-linux/flash_attn_interface.py ADDED Viewed

	@@ -0,0 +1,1708 @@

+# Copyright (c) 2023, Tri Dao.
+from typing import Optional, Sequence, Tuple, Union
+import torch
+import torch.nn as nn
+import os
+# # isort: off
+# # We need to import the CUDA kernels after importing torch
+# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE"
+# if USE_TRITON_ROCM:
+#     from .flash_attn_triton_amd import interface_fa as flash_attn
+# else:
+#     import flash_attn_2_cuda as flash_attn
+from ._ops import ops as flash_attn
+from ._ops import add_op_namespace_prefix
+# # isort: on
+# Torch >= 2.13 can pass a Generator through the stable ABI. Older versions reject the argument
+# before the kernel runs, so there the CUDA op has to draw its Philox seed and offset itself.
+_STABLE_ABI_GENERATOR = torch.__version__ >= "2.13"
+def maybe_contiguous(x):
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+def _get_device():
+    if torch.xpu.is_available():
+        return "xpu"
+    elif torch.cuda.is_available():
+        return "cuda"
+    else:
+        return "cpu"
+def _get_block_size_n(device, head_dim, is_dropout, is_causal):
+    assert head_dim <= 256
+    if device.type == "xpu":
+        if head_dim <= 96:
+            return 64
+        elif head_dim <= 128:
+            return 32
+        elif head_dim <= 256:
+            return 64
+        else:
+            return 32
+    # This should match the block sizes in the CUDA kernel
+    major, minor = torch.cuda.get_device_capability(device)
+    is_sm8x = (
+        major == 8 and minor > 0
+    )  # Only include sm86 and sm89, exclude sm80 (A100)
+    is_sm80 = major == 8 and minor == 0
+    is_sm90 = major == 9 and minor == 0
+    if head_dim <= 32:
+        return 128
+    if head_dim <= 64:
+        return 128 if not is_dropout else 64
+    elif head_dim <= 96:
+        return 64
+    elif head_dim <= 128:
+        if is_sm8x:
+            return 64 if (not is_dropout and is_causal) else 32
+        else:
+            return 64 if not is_dropout else 32
+    elif head_dim <= 192:
+        return 64
+    elif head_dim <= 224:
+        return 64
+    elif head_dim <= 256:
+        return 64
+def round_multiple(x, m):
+    return (x + m - 1) // m * m
+def _philox_generator(
+    q: torch.Tensor, dropout_p: float, generator: Optional[torch.Generator] = None
+) -> Optional[torch.Generator]:
+    """The generator to hand to `fwd` / `varlen_fwd`.
+    With dropout on CUDA this is the default generator unless one is given, so the kernel reserves
+    its seed and offset exactly like the non-stable build (see set_params_philox in flash_api.cpp).
+    """
+    if not q.is_cuda:
+        return generator
+    if dropout_p <= 0:
+        return None
+    if not _STABLE_ABI_GENERATOR:
+        if generator is not None:
+            raise ValueError(
+                "flash_attn2: passing a generator to the CUDA kernel requires torch >= 2.13; "
+                "seed the default generator with torch.manual_seed instead"
+            )
+        return None
+    if generator is None:
+        generator = torch.cuda.default_generators[q.device.index]
+    return generator
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_forward"),
+    mutates_args=(),
+    device_types=_get_device(),
+    # Tells torch.compile the forward draws random numbers, so it is never deduplicated or
+    # recomputed (e.g. under activation checkpointing).
+    tags=(torch.Tag.nondeterministic_seeded,),
+)
+def _flash_attn_forward(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    return_softmax: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    gen = _philox_generator(q, dropout_p)
+    out, softmax_lse, S_dmask, rng_state = flash_attn.fwd(
+        q,
+        k,
+        v,
+        None,
+        alibi_slopes,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out, softmax_lse, S_dmask, rng_state
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward"))
+def _flash_attn_forward_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    return_softmax: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    batch_size, seqlen_q, num_heads, head_size = q.shape
+    seqlen_k = k.shape[1]
+    out = torch.empty_like(q)
+    softmax_lse = torch.empty(
+        (batch_size, num_heads, seqlen_q),
+        dtype=torch.float32,
+        device=q.device,
+        layout=q.layout,
+    )
+    p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout)
+    if return_softmax:
+        p = torch.empty(
+            (
+                batch_size,
+                num_heads,
+                round_multiple(seqlen_q, 128),
+                round_multiple(seqlen_k, 128),
+            ),
+            dtype=q.dtype,
+            device=q.device,
+            layout=q.layout,
+        )
+    rng_state = torch.empty((2,), dtype=torch.int64, device=q.device)
+    return out, softmax_lse, p, rng_state
+_wrapped_flash_attn_forward = flash_attn._flash_attn_forward  # upstream alias, see CARD.md
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_varlen_forward"),
+    mutates_args=(),
+    device_types=_get_device(),
+    tags=(torch.Tag.nondeterministic_seeded,),
+)
+def _flash_attn_varlen_forward(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    return_softmax: bool = False,
+    block_table: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    gen = _philox_generator(q, dropout_p)
+    out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd(
+        q,
+        k,
+        v,
+        None,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_k,
+        leftpad_k,
+        block_table,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        zero_tensors,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    # if out.isnan().any() or softmax_lse.isnan().any():
+    #     breakpoint()
+    return out, softmax_lse, S_dmask, rng_state
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_forward"))
+def _flash_attn_varlen_forward_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    return_softmax: bool = False,
+    block_table: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    paged_kv = block_table is not None
+    batch_size = cu_seqlens_q.numel() - 1
+    total_q, num_heads, _ = q.shape
+    out = torch.empty_like(q)
+    softmax_lse = torch.empty(
+        (num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout
+    )
+    p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout)
+    seqlen_q_rounded = round_multiple(max_seqlen_q, 128)
+    seqlen_k_rounded = round_multiple(max_seqlen_k, 128)
+    if return_softmax:
+        p = torch.empty(
+            (batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded),
+            dtype=q.dtype,
+            device=q.device,
+            layout=q.layout,
+        )
+    rng_state = torch.empty((2,), dtype=torch.int64, device=q.device)
+    return out, softmax_lse, p, rng_state
+_wrapped_flash_attn_varlen_forward = flash_attn._flash_attn_varlen_forward  # upstream alias, see CARD.md
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_backward"),
+    mutates_args=("dq", "dk", "dv"),
+    device_types=_get_device(),
+)
+def _flash_attn_backward(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    # dq, dk, dv are allocated by us so they should already be contiguous
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    (
+        dq,
+        dk,
+        dv,
+        softmax_d,
+    ) = flash_attn.bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        alibi_slopes,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        None,
+        rng_state,
+    )
+    return softmax_d
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward"))
+def _flash_attn_backward_fake(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    if dq is None:
+        dq = torch.empty_like(q)
+    if dk is None:
+        dk = torch.empty_like(k)
+    if dv is None:
+        dv = torch.empty_like(v)
+    batch_size, seqlen_q, num_heads, _ = q.shape
+    softmax_d = torch.empty(
+        (batch_size, num_heads, round_multiple(seqlen_q, 128)),
+        device=q.device,
+        dtype=torch.float32,
+    )
+    return softmax_d
+_wrapped_flash_attn_backward = flash_attn._flash_attn_backward  # upstream alias, see CARD.md
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_varlen_backward"),
+    mutates_args=("dq", "dk", "dv"),
+    device_types=_get_device(),
+)
+def _flash_attn_varlen_backward(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> torch.Tensor:
+    # dq, dk, dv are allocated by us so they should already be contiguous
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    (
+        dq,
+        dk,
+        dv,
+        softmax_d,
+    ) = flash_attn.varlen_bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        zero_tensors,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        None,
+        rng_state,
+    )
+    # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any():
+    #     breakpoint()
+    return softmax_d
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_backward"))
+def _flash_attn_varlen_backward_fake(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> torch.Tensor:
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    batch_size = cu_seqlens_q.numel() - 1
+    total_q, num_heads, _ = q.shape
+    if dq is None:
+        dq = torch.empty_like(q)
+    if dk is None:
+        dk = torch.empty_like(k)
+    if dv is None:
+        dv = torch.empty_like(v)
+    softmax_d = torch.empty(
+        (num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32
+    )
+    return softmax_d
+_wrapped_flash_attn_varlen_backward = flash_attn._flash_attn_varlen_backward  # upstream alias, see CARD.md
+class FlashAttnQKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        qkv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and qkv.requires_grad
+        if softmax_scale is None:
+            softmax_scale = qkv.shape[-1] ** (-0.5)
+        q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach()
+        head_size_og = q.size(3)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+        )
+        if is_grad:
+            ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
+        qkv_shape = q.shape[:-2] + (3, *q.shape[-2:])
+        dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device)
+        head_size_og = dout.size(3)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dqkv[:, :, 0],
+            dqkv[:, :, 1],
+            dqkv[:, :, 2],
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dqkv = dqkv[..., : dout.shape[-1]]  # We could have padded the head dimension
+        return dqkv, None, None, None, None, None, None, None, None, None
+class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        qkv,
+        cu_seqlens,
+        max_seqlen,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and qkv.requires_grad
+        if softmax_scale is None:
+            softmax_scale = qkv.shape[-1] ** (-0.5)
+        q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach()
+        head_size_og = q.size(2)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
+            q,
+            k,
+            v,
+            cu_seqlens,
+            cu_seqlens,
+            max_seqlen,
+            max_seqlen,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+            block_table=None,
+        )
+        if is_grad:
+            ctx.save_for_backward(
+                q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state
+            )
+            ctx.dropout_p = dropout_p
+            ctx.max_seqlen = max_seqlen
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors
+        qkv_shape = q.shape[:-2] + (3, *q.shape[-2:])
+        dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device)
+        head_size_og = dout.size(2)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_varlen_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dqkv[:, 0],
+            dqkv[:, 1],
+            dqkv[:, 2],
+            cu_seqlens,
+            cu_seqlens,
+            ctx.max_seqlen,
+            ctx.max_seqlen,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dqkv = dqkv[..., : dout.shape[-1]]  # We could have padded the head dimension
+        return dqkv, None, None, None, None, None, None, None, None, None, None, None
+class FlashAttnKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        kv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach()
+        head_size_og = q.size(3)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+        )
+        if is_grad:
+            ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
+        dq = torch.empty_like(q)
+        kv_shape = k.shape[:-2] + (2, *k.shape[-2:])
+        dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device)
+        head_size_og = dout.size(3)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dkv[:, :, 0],
+            dkv[:, :, 1],
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dkv = dkv[..., : dout.shape[-1]]
+        return dq, dkv, None, None, None, None, None, None, None, None, None
+class FlashAttnVarlenKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        kv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        k, v = kv[:, 0].detach(), kv[:, 1].detach()
+        head_size_og = q.size(2)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+            block_table=None,
+        )
+        if is_grad:
+            ctx.save_for_backward(
+                q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state
+            )
+            ctx.dropout_p = dropout_p
+            ctx.max_seqlen_q = max_seqlen_q
+            ctx.max_seqlen_k = max_seqlen_k
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = (
+            ctx.saved_tensors
+        )
+        dq = torch.empty_like(q)
+        kv_shape = k.shape[:-2] + (2, *k.shape[-2:])
+        dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device)
+        head_size_og = dout.size(2)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_varlen_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dkv[:, 0],
+            dkv[:, 1],
+            cu_seqlens_q,
+            cu_seqlens_k,
+            ctx.max_seqlen_q,
+            ctx.max_seqlen_k,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dkv = dkv[..., : dout.shape[-1]]
+        return (
+            dq,
+            dkv,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+class FlashAttnFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        head_size_og = q.size(3)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+        )
+        if is_grad:
+            ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
+        dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+        head_size_og = dout.size(3)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dk,
+            dv,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dk = dk[..., : dout.shape[-1]]
+        dv = dv[..., : dout.shape[-1]]
+        return dq, dk, dv, None, None, None, None, None, None, None, None, None
+class FlashAttnVarlenFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        block_table,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        head_size_og = q.size(2)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+            block_table=block_table,
+        )
+        if is_grad:
+            ctx.save_for_backward(
+                q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state
+            )
+            ctx.dropout_p = dropout_p
+            ctx.max_seqlen_q = max_seqlen_q
+            ctx.max_seqlen_k = max_seqlen_k
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = (
+            ctx.saved_tensors
+        )
+        dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+        head_size_og = dout.size(2)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_varlen_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dk,
+            dv,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            ctx.max_seqlen_q,
+            ctx.max_seqlen_k,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dk = dk[..., : dout.shape[-1]]
+        dv = dv[..., : dout.shape[-1]]
+        return (
+            dq,
+            dk,
+            dv,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+def flash_attn_qkvpacked_func(
+    qkv,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # <=0.0 means deactivate
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If Q, K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of Q, K, V.
+    For multi-query and grouped-query attention (MQA/GQA), please see
+    flash_attn_kvpacked_func and flash_attn_func.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.
+    Arguments:
+        qkv: (batch_size, seqlen, 3, nheads, headdim)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to
+            the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnQKVPackedFunc.apply(
+        qkv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_kvpacked_func(
+    q,
+    kv,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of K, V.
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        kv: (batch_size, seqlen, 2, nheads_k, headdim)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnKVPackedFunc.apply(
+        q,
+        kv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_func(
+    q,
+    k,
+    v,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        k: (batch_size, seqlen, nheads_k, headdim)
+        v: (batch_size, seqlen, nheads_k, headdim)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnFunc.apply(
+        q,
+        k,
+        v,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_varlen_qkvpacked_func(
+    qkv,
+    cu_seqlens,
+    max_seqlen,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If Q, K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_varlen_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of Q, K, V.
+    For multi-query and grouped-query attention (MQA/GQA), please see
+    flash_attn_varlen_kvpacked_func and flash_attn_varlen_func.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.
+    Arguments:
+        qkv: (total, 3, nheads, headdim), where total = total number of tokens in the batch.
+        cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into qkv.
+        max_seqlen: int. Maximum sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnVarlenQKVPackedFunc.apply(
+        qkv,
+        cu_seqlens,
+        max_seqlen,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_varlen_kvpacked_func(
+    q,
+    kv,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of K, V.
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.
+        kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into q.
+        cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into kv.
+        max_seqlen_q: int. Maximum query sequence length in the batch.
+        max_seqlen_k: int. Maximum key sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnVarlenKVPackedFunc.apply(
+        q,
+        kv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_varlen_func(
+    q,
+    k,
+    v,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+    block_table=None,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.
+        k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into q.
+        cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into kv.
+        max_seqlen_q: int. Maximum query sequence length in the batch.
+        max_seqlen_k: int. Maximum key sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnVarlenFunc.apply(
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        block_table,
+        False if q.device.type == "cpu" else torch.is_grad_enabled(),
+    )
+def flash_attn_with_kvcache(
+    q,
+    k_cache,
+    v_cache,
+    k=None,
+    v=None,
+    rotary_cos=None,
+    rotary_sin=None,
+    cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None,
+    cache_batch_idx: Optional[torch.Tensor] = None,
+    cache_leftpad: Optional[torch.Tensor] = None,
+    block_table: Optional[torch.Tensor] = None,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    rotary_interleaved=True,
+    alibi_slopes=None,
+    num_splits=0,
+    return_softmax_lse=False,
+):
+    """
+    If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from
+    k and v. This is useful for incremental decoding: you can pass in the cached keys/values from
+    the previous step, and update them with the new keys/values from the current step, and do
+    attention with the updated cache, all in 1 kernel.
+    If you pass in k / v, you must make sure that the cache is large enough to hold the new values.
+    For example, the KV cache could be pre-allocated with the max sequence length, and you can use
+    cache_seqlens to keep track of the current sequence lengths of each sequence in the batch.
+    Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be
+    rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
+    If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos
+    and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
+    If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at
+    indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens).
+    See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function.
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Note: Does not support backward pass.
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,
+            or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)
+            page_block_size must be a multiple of 256.
+        v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,
+            or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)
+        k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate
+            k with k_cache, starting at the indices specified by cache_seqlens.
+        v [optional]: (batch_size, seqlen_new, nheads_k, headdim). Similar to k.
+        rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding
+            to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16.
+        rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos.
+        cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the
+            KV cache.
+        cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache.
+            If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1].
+            If the indices are not distinct, and k and v are provided, the values updated in the cache
+                 might come from any of the duplicate indices.
+        cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0.
+        block_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in.
+            If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False,
+            rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1
+            (i.e. GPT-NeoX style).
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        num_splits: int. If > 1, split the key/value into this many chunks along the sequence.
+           If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic
+           to automatically determine the number of splits.
+           Don't change this unless you know what you are doing.
+        return_softmax_lse: bool. Whether to return the logsumexp of the attention scores.
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+    """
+    assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension"
+    assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension"
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** (-0.5)
+    if cache_seqlens is not None and isinstance(cache_seqlens, int):
+        cache_seqlens = torch.full(
+            (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device
+        )
+        cache_seqlens = maybe_contiguous(cache_seqlens)
+    cache_batch_idx = maybe_contiguous(cache_batch_idx)
+    block_table = maybe_contiguous(block_table)
+    out, softmax_lse = flash_attn.fwd_kvcache(
+        q,
+        k_cache,
+        v_cache,
+        k,
+        v,
+        cache_seqlens,
+        rotary_cos,
+        rotary_sin,
+        cache_batch_idx,
+        cache_leftpad,
+        block_table,
+        alibi_slopes,
+        None,
+        softmax_scale,
+        causal,
+        window_size[0],
+        window_size[1],
+        softcap,
+        rotary_interleaved,
+        num_splits,
+    )
+    return (out, softmax_lse) if return_softmax_lse else out

build/torch-stable-abi210-cu128-aarch64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,39 @@

+{
+  "name": "flash-attn2",
+  "id": "_flash_attn2_cuda_460047c",
+  "version": 4,
+  "kernels-minver": "0.14.0",
+  "license": "BSD-3-Clause",
+  "python-depends": [],
+  "kernel-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Nz8hqrusjoUHQqszpxfF2USMwUXn1VfpvcmtAAvutSM=",
+      "_flash_attn2_cuda_460047c.abi3.so": "y2DTM0Gg8SYvoR2RnAd2OYZR4yfN/rfSH608Kur4UTs=",
+      "_ops.py": "DV4Eo4+s3aaNmDsfXoZiYKNLnRUJqlsdDGlMsdKFgk4=",
+      "flash_attn_interface.py": "6UYUhfUax4+yOUygyS7K8dVjrIjU5GQr1AJ2gSp8rPo=",
+      "symbols.json": "uvC+87tHUhtHv7L/Nl53T/Ms7ceX317/jwQBoqHY14A="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.18.0-dev0",
+      "commit": "d054906a13710b811b96697b4635f03020842733",
+      "dirty": false
+    },
+    "kernel": {
+      "commit": "460047c291fc7f0b3dbd4239b46bd5e1281b1806",
+      "dirty": false
+    }
+  }
+}

build/torch-stable-abi210-cu128-aarch64-linux/metadata.json.sigstore ADDED Viewed

	@@ -0,0 +1 @@

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHKzCCBrGgAwIBAgIUGbk1nmy8ZtAxTYtkF97vt1SsF5wwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDI0MDAwNzQ4WhcNMjYxMDI0MDAxNzQ4WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE4WV6kHnDLUJYxjhxaSnJwr6NG2lryfV2sMRL2KY8sReoSRkNxXbSD6QrvWDaIY/r8pOm0grvRwRB7QZvafDFoqOCBdAwggXMMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBCg0NjAwNDdjMjkxZmM3ZjBiM2RiZDQyMzliNDZiZDVlMTI4MWIxODA2MBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDCg0NjAwNDdjMjkxZmM3ZjBiM2RiZDQyMzliNDZiZDVlMTI4MWIxODA2MBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDCg0NjAwNDdjMjkxZmM3ZjBiM2RiZDQyMzliNDZiZDVlMTI4MWIxODA2MB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoNDYwMDQ3YzI5MWZjN2YwYjNkYmQ0MjM5YjQ2YmQ1ZTEyODFiMTgwNjAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTk4MjY0MDUxNS9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGJBgorBgEEAdZ5AgQCBHsEeQB3AHUA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhazwspAAABAMARjBEAiAcFNKlbzabb93VQ+7byyb/FyWB5TroTsHwjqvnzTK8nAIgW5ZevxlbnBmf4yzzuXH0Oja6pknFoj1ZKRUFNpnv8eEwCgYIKoZIzj0EAwMDaAAwZQIxAOa53j1w0YZOSPrUzRCUOXQbvkzrt5KikVkbWyWTPga4xZOAuvSGjqij1kWeLEuozQIwOp9POt2v6NKJ7zOdkI0fVHPXEw19zksIee8Evng7rIfy8H48SPQboIo0hpmOnEhD"},"tlogEntries":[{"logIndex":"3430231150","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792800468","inclusionPromise":{"signedEntryTimestamp":"MEQCIGZLwd3J3l3V5DJXhWg+9EGadjCvaCWzPMsUG93F6kNJAiAP+kvJFeyEBb1JoFaYPSRaVko7A3De8lLTG9MmgfItug=="},"inclusionProof":{"logIndex":"3308326888","rootHash":"0dxDHKHJJNB18GUn1mFkbYi+mmMHAiKBJZiSnPoSq0M=","treeSize":"3308327082","hashes":["tEmdr9KY4X4wGUqr7YjGpLBDSzBdktM+iTNHlN/KVnI=","28a1JPd0NfARI9ataGIjKR2HImyILikAgZQlo9g3gR0=","HPdmEIkQK6TAzqM3AhOjMT2pV8lJmiBdHB4ecJvoH3c=","1oNpBskEERK7M/2IelxCy6dVBEIOAlrvMrekGwOEIhQ=","4ihQi/htRBQaP0gx9I9zFdpjlhs98HBFIi6y7MxxGEs=","fhi0wkcZ9wIt9C3ZoYSTiwPoztVF0FoG/M0C7cRQIjs=","/NfyecAkSIGEAZki6MTEL1vFqGHHJqwpejFj36u/6W8=","rIuKGGpr2xHVfYpBBBRmjjC1BSLVSY7CUZBWo+YwVe4=","D4KiQ6dLka04xW5SItSpf6u70l4OLIJY9TGL3Vgb7J4=","JHbO9SM0uiPdIUJoCFWlrHv/QNK/f6UMlwBq/6914+w=","PuwdRME/93l6s2BMRjCojmjsVC29KnBqdtPSPE1l6cQ=","tWdxtMrYqmZ0zKHnOCAERPP7SvuDvubtw3jwj7rH4uA=","kM+kX6uNUuDAQM/IB5IDAgZEYbCs+QGUBYxh34/qyC4=","8x0N/oAmvpY+6DQUGTU+i6a1Cuqphro2WbV5WZ/RHL8=","DMchGrSAeKV9qHo75XJV3UFZ/CHF8zofPRu+NgVRwA4=","V6dXjeDcwYxBG/9gVSBM8rYdPrAqGpZsf1e7v6l67GA=","iDfImp65w8dirRfYGNgFA34JTUYDWbcJwNBBU/Ok+X4=","YYCr61A8azar3dfO8Ct+pri7LqDxkUlUkfs7fiqYErU=","1DaDKzrCNNBhmSdY5vFjz+NpAmxmYfYEgBCdPLrMP1A=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3308327082\n0dxDHKHJJNB18GUn1mFkbYi+mmMHAiKBJZiSnPoSq0M=\n\n— rekor.sigstore.dev wNI9ajBFAiBW2C60YRvrM9xE+1c/VFt32boXOn4+Am6CokqDb2PTJAIhALsoOhB00Myr6SMsAdAtJZkX1K3OPOb9DU6y/28bhshI\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJjNjRiNmQxNWQ4NWNkZTQ5YzJmMjQ1NTJiMmNiN2MyY2QyNWE3MTE5ZDlkOTRjOTg5OTQzMWNkOTQxNWRlNzFjIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUQ4aGFKZnFUSXkrSVVuRWxaK3NiM0tNWFRkUUF0OVF3K3grNENTOXI0MklBSWdFbXZlalF2dXBEVjd1UzFsYjNhUWxkcE9HNG5VeElmMmFIc0F6MlZOd0NVPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaExla05EUW5KSFowRjNTVUpCWjBsVlIySnJNVzV0ZVRoYWRFRjRWRmwwYTBZNU4zWjBNVk56UmpWM2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFU1RCTlJFRjNUbnBSTkZkb1kwNU5hbGw0VFVSSk1FMUVRWGhPZWxFMFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVUwVjFZMmEwaHVSRXhWU2xsNGFtaDRZVk51U25keU5rNUhNbXh5ZVdaV01uTk5Va3dLTWt0Wk9ITlNaVzlUVW10T2VGaGlVMFEyVVhKMlYwUmhTVmt2Y2pod1QyMHdaM0oyVW5kU1FqZFJXblpoWmtSR2IzRlBRMEprUVhkbloxaE5UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmN3VG1wQmR3cE9SR1JxVFdwcmVGcHRUVE5hYWtKcFRUSlNhVnBFVVhsTmVteHBUa1JhYVZwRVZteE5WRWswVFZkSmVFOUVRVEpOUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEWnpCT2FrRjNUa1JrYWsxcWEzaGFiVTB6Q2xwcVFtbE5NbEpwV2tSUmVVMTZiR2xPUkZwcFdrUldiRTFVU1RSTlYwbDRUMFJCTWsxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEWnpCT2FrRjNDazVFWkdwTmFtdDRXbTFOTTFwcVFtbE5NbEpwV2tSUmVVMTZiR2xPUkZwcFdrUldiRTFVU1RSTlYwbDRUMFJCTWsxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlPUkZsM1RVUlJNMWw2U1RWTlYxcHFUakpaZDFscVRtdFpiVkV3VFdwTk5WbHFVVEpaYlZFeFdsUkZlVTlFUm1rS1RWUm5kMDVxUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUnJORTFxV1RCTlJGVjRUbE01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkS0NrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NITkZaVkZDTTBGSVZVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvWVhwM2MzQkJRVUZDUVUxQlVtcENSVUZwUVdOR1RrdHNZbnBoWW1JNU0xWlJLemRpZVhsaUwwWjVWMEkxVkhKdlZITklkd3BxY1hadWVsUkxPRzVCU1dkWE5WcGxkbmhzWW01Q2JXWTBlWHA2ZFZoSU1FOXFZVFp3YTI1R2Iyb3hXa3RTVlVaT2NHNTJPR1ZGZDBObldVbExiMXBKQ25wcU1FVkJkMDFFWVVGQmQxcFJTWGhCVDJFMU0yb3hkekJaV2s5VFVISlZlbEpEVlU5WVVXSjJhM3B5ZERWTGFXdFdhMkpYZVZkVVVHZGhOSGhhVDBFS2RYWlRSMnB4YVdveGExZGxURVYxYjNwUlNYZFBjRGxRVDNReWRqWk9TMG8zZWs5a2Ewa3dabFpJVUZoRmR6RTVlbXR6U1dWbE9FVjJibWMzY2tsbWVRbzRTRFE0VTFCUlltOUpiekJvY0cxUGJrVm9SQW90TFMwdExVVk9SQ0JEUlZKVVNVWkpRMEZVUlMwdExTMHRDZz09In19fX0="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgJQV5D23yDAIgDTv4Yv16z9aMxggzPuqToqoPm4v8O38CFEC/mLVnKnZKvq7p+4Od4Odmb5WGGA8yMDI2MTAyNDAwMDc0OFowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDI0MDAwNzQ4WjAvBgkqhkiG9w0BCQQxIgQgenxrO7f3Yd1WfiHFvp5639SGocPIHs05Wki1wZGRHSQwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIxAImlABvKBRPfu11mLxuiccH4LMnQ1qpM29/zxDdkh2Pnz0Kprgm0GSFMmudHlfC1fwIwHfUuF271l/bKNDbCMzTU4D82yvTYXe3RZTPvSXZN54VDTyO8WXA9uj5Pmd4Yuf3B"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"xkttFdhc3knC8kVSsst8LNJacRnZ2UyYmUMc2UFd5xw="},"signature":"MEUCIQD8haJfqTIy+IUnElZ+sb3KMXTdQAt9Qw+x+4CS9r42IAIgEmvejQvupDV7uS1lb3aQldpOG4nUxIf2aHsAz2VNwCU="}}

build/torch-stable-abi210-cu128-aarch64-linux/symbols.json ADDED Viewed

	@@ -0,0 +1,1273 @@

+{
+  "schema_version": 1,
+  "module": "_flash_attn2_cuda_460047c",
+  "functions": [
+    {
+      "name": "flash_attn_func",
+      "qualname": "flash_attn_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (batch_size, seqlen, nheads, headdim)\n    k: (batch_size, seqlen, nheads_k, headdim)\n    v: (batch_size, seqlen, nheads_k, headdim)\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_kvpacked_func",
+      "qualname": "flash_attn_kvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of K, V.\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (batch_size, seqlen, nheads, headdim)\n    kv: (batch_size, seqlen, 2, nheads_k, headdim)\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "kv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_qkvpacked_func",
+      "qualname": "flash_attn_qkvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf Q, K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of Q, K, V.\nFor multi-query and grouped-query attention (MQA/GQA), please see\nflash_attn_kvpacked_func and flash_attn_func.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.\n\nArguments:\n    qkv: (batch_size, seqlen, 3, nheads, headdim)\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to\n        the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "qkv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_varlen_func",
+      "qualname": "flash_attn_varlen_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.\n    k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.\n    v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.\n    cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into q.\n    cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into kv.\n    max_seqlen_q: int. Maximum query sequence length in the batch.\n    max_seqlen_k: int. Maximum key sequence length in the batch.\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (total, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_varlen_kvpacked_func",
+      "qualname": "flash_attn_varlen_kvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of K, V.\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.\n    kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch.\n    cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into q.\n    cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into kv.\n    max_seqlen_q: int. Maximum query sequence length in the batch.\n    max_seqlen_k: int. Maximum key sequence length in the batch.\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (total, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "kv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_varlen_qkvpacked_func",
+      "qualname": "flash_attn_varlen_qkvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf Q, K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_varlen_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of Q, K, V.\nFor multi-query and grouped-query attention (MQA/GQA), please see\nflash_attn_varlen_kvpacked_func and flash_attn_varlen_func.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.\n\nArguments:\n    qkv: (total, 3, nheads, headdim), where total = total number of tokens in the batch.\n    cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into qkv.\n    max_seqlen: int. Maximum sequence length in the batch.\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (total, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "qkv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_with_kvcache",
+      "qualname": "flash_attn_with_kvcache",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from\nk and v. This is useful for incremental decoding: you can pass in the cached keys/values from\nthe previous step, and update them with the new keys/values from the current step, and do\nattention with the updated cache, all in 1 kernel.\n\nIf you pass in k / v, you must make sure that the cache is large enough to hold the new values.\nFor example, the KV cache could be pre-allocated with the max sequence length, and you can use\ncache_seqlens to keep track of the current sequence lengths of each sequence in the batch.\n\nAlso apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be\nrotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.\nIf causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos\nand rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.\nIf not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at\nindices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens).\n\nSee tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function.\n\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nNote: Does not support backward pass.\n\nArguments:\n    q: (batch_size, seqlen, nheads, headdim)\n    k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,\n        or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)\n        page_block_size must be a multiple of 256.\n    v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,\n        or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)\n    k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate\n        k with k_cache, starting at the indices specified by cache_seqlens.\n    v [optional]: (batch_size, seqlen_new, nheads_k, headdim). Similar to k.\n    rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding\n        to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16.\n    rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos.\n    cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the\n        KV cache.\n    cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache.\n        If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1].\n        If the indices are not distinct, and k and v are provided, the values updated in the cache\n             might come from any of the duplicate indices.\n    cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0.\n    block_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in.\n        If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False,\n        rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1\n        (i.e. GPT-NeoX style).\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    num_splits: int. If > 1, split the key/value into this many chunks along the sequence.\n       If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic\n       to automatically determine the number of splits.\n       Don't change this unless you know what you are doing.\n    return_softmax_lse: bool. Whether to return the logsumexp of the attention scores.\n\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k_cache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "v_cache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "rotary_cos",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "rotary_sin",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "cache_seqlens",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "int | torch.Tensor | None"
+          },
+          {
+            "name": "cache_batch_idx",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "cache_leftpad",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "rotary_interleaved",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "True",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "num_splits",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": null
+          },
+          {
+            "name": "return_softmax_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "fwd",
+      "qualname": "fwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Forward pass for multi-head attention.\n\nArgs:\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    out: Optional output tensor, same shape as q\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    return_softmax: Whether to return softmax weights\n    gen: Optional random number generator\n\nReturns:\n    List of tensors: [output, softmax_lse, (softmax if return_softmax)]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "return_softmax",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "varlen_fwd",
+      "qualname": "varlen_fwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Forward pass for multi-head attention with variable sequence lengths.\n\nArgs:\n    q: Query tensor of shape [total_q, num_heads, head_size]\n    k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]\n    cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]\n    out: Optional output tensor of shape [total_q, num_heads, head_size]\n    seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size]\n    leftpad_k: Optional left padding for keys of shape [batch_size]\n    block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    max_seqlen_q: Maximum sequence length for queries\n    max_seqlen_k: Maximum sequence length for keys\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    zero_tensors: Whether to zero tensors before computation\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    return_softmax: Whether to return softmax weights\n    gen: Optional random number generator\n\nReturns:\n    List of tensors: [output, softmax_lse, (softmax if return_softmax)]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "seqused_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "leftpad_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "zero_tensors",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "return_softmax",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "bwd",
+      "qualname": "bwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Backward pass for multi-head attention.\n\nArgs:\n    dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]\n    softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]\n    dq: Optional gradient tensor for queries, same shape as q\n    dk: Optional gradient tensor for keys, same shape as k\n    dv: Optional gradient tensor for values, same shape as v\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    deterministic: Whether to use deterministic algorithms\n    gen: Optional random number generator\n    rng_state: Optional RNG state from forward pass\n\nReturns:\n    List of tensors: [dq, dk, dv]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "dout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "softmax_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "dq",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dk",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          },
+          {
+            "name": "rng_state",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "varlen_bwd",
+      "qualname": "varlen_bwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Backward pass for multi-head attention with variable sequence lengths.\n\nArgs:\n    dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]\n    softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]\n    cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]\n    cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]\n    dq: Optional gradient tensor for queries, same shape as q\n    dk: Optional gradient tensor for keys, same shape as k\n    dv: Optional gradient tensor for values, same shape as v\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    max_seqlen_q: Maximum sequence length for queries\n    max_seqlen_k: Maximum sequence length for keys\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    zero_tensors: Whether to zero tensors before computation\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    deterministic: Whether to use deterministic algorithms\n    gen: Optional random number generator\n    rng_state: Optional RNG state from forward pass\n\nReturns:\n    List of tensors: [dq, dk, dv]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "dout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "softmax_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "dq",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dk",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "zero_tensors",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          },
+          {
+            "name": "rng_state",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "fwd_kvcache",
+      "qualname": "fwd_kvcache",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Forward pass for multi-head attention with KV cache.\n\nArgs:\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]\n    v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]\n    seqlens_k: Optional sequence lengths for keys of shape [batch_size]\n    rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2]\n    rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2]\n    cache_batch_idx: Optional indices to index into the KV cache\n    leftpad_k: Optional left padding for keys of shape [batch_size]\n    block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    out: Optional output tensor, same shape as q\n    softmax_scale: Scale factor for softmax\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    is_rotary_interleaved: Whether rotary embeddings are interleaved\n    num_splits: Number of splits for computation\n\nReturns:\n    List of tensors: [output, softmax_lse]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "kcache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "vcache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "rotary_cos",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "rotary_sin",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "cache_batch_idx",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "leftpad_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "is_rotary_interleaved",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "num_splits",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "1",
+            "annotation": "int"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    }
+  ],
+  "layers": []
+}

build/torch-stable-abi210-cu130-aarch64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,424 @@

+from typing import List, Optional
+import torch
+from . import flash_attn_interface
+from ._ops import ops as flash_attn_ops
+from .flash_attn_interface import (
+    # helper to set up the Philox random number generator for dropout
+    _philox_generator,
+    # reexported for convenience to be used in the high-level API functions
+    flash_attn_func,
+    flash_attn_kvpacked_func,
+    flash_attn_qkvpacked_func,
+    flash_attn_varlen_func,
+    flash_attn_varlen_kvpacked_func,
+    flash_attn_varlen_qkvpacked_func,
+    flash_attn_with_kvcache,
+)
+__all__ = [
+    # High-level API (with autograd support)
+    "flash_attn_func",
+    "flash_attn_kvpacked_func",
+    "flash_attn_qkvpacked_func",
+    "flash_attn_varlen_func",
+    "flash_attn_varlen_kvpacked_func",
+    "flash_attn_varlen_qkvpacked_func",
+    "flash_attn_with_kvcache",
+    # Upstream interface module. Its `_flash_attn_forward`, `_flash_attn_backward`,
+    # `_flash_attn_varlen_forward` and `_flash_attn_varlen_backward` are supported
+    # low-level entry points with upstream's signatures: Torch custom ops without
+    # autograd that return the LSE and are traceable by torch.compile. The autograd
+    # Functions above are built on top of them.
+    "flash_attn_interface",
+    # Low-level ops (no autograd)
+    "fwd",
+    "varlen_fwd",
+    "bwd",
+    "varlen_bwd",
+    "fwd_kvcache",
+]
+def fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    return_softmax: bool = False,
+    gen: Optional[torch.Generator] = None,
+) -> List[torch.Tensor]:
+    """
+    Forward pass for multi-head attention.
+    Args:
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        out: Optional output tensor, same shape as q
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        return_softmax: Whether to return softmax weights
+        gen: Optional random number generator
+    Returns:
+        List of tensors: [output, softmax_lse, (softmax if return_softmax)]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    gen = _philox_generator(q, p_dropout, gen)
+    return flash_attn_ops.fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+def varlen_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    block_table: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    max_seqlen_q: int = 0,
+    max_seqlen_k: int = 0,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    zero_tensors: bool = False,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    return_softmax: bool = False,
+    gen: Optional[torch.Generator] = None,
+) -> List[torch.Tensor]:
+    """
+    Forward pass for multi-head attention with variable sequence lengths.
+    Args:
+        q: Query tensor of shape [total_q, num_heads, head_size]
+        k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]
+        cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]
+        out: Optional output tensor of shape [total_q, num_heads, head_size]
+        seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size]
+        leftpad_k: Optional left padding for keys of shape [batch_size]
+        block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        max_seqlen_q: Maximum sequence length for queries
+        max_seqlen_k: Maximum sequence length for keys
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        zero_tensors: Whether to zero tensors before computation
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        return_softmax: Whether to return softmax weights
+        gen: Optional random number generator
+    Returns:
+        List of tensors: [output, softmax_lse, (softmax if return_softmax)]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    gen = _philox_generator(q, p_dropout, gen)
+    return flash_attn_ops.varlen_fwd(
+        q,
+        k,
+        v,
+        out,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_k,
+        leftpad_k,
+        block_table,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        p_dropout,
+        softmax_scale,
+        zero_tensors,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+def bwd(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor] = None,
+    dk: Optional[torch.Tensor] = None,
+    dv: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    deterministic: bool = False,
+    gen: Optional[torch.Generator] = None,
+    rng_state: Optional[torch.Tensor] = None,
+) -> List[torch.Tensor]:
+    """
+    Backward pass for multi-head attention.
+    Args:
+        dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]
+        softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]
+        dq: Optional gradient tensor for queries, same shape as q
+        dk: Optional gradient tensor for keys, same shape as k
+        dv: Optional gradient tensor for values, same shape as v
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        deterministic: Whether to use deterministic algorithms
+        gen: Optional random number generator
+        rng_state: Optional RNG state from forward pass
+    Returns:
+        List of tensors: [dq, dk, dv]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    return flash_attn_ops.bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        gen,
+        rng_state,
+    )
+def varlen_bwd(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    dq: Optional[torch.Tensor] = None,
+    dk: Optional[torch.Tensor] = None,
+    dv: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    max_seqlen_q: int = 0,
+    max_seqlen_k: int = 0,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    zero_tensors: bool = False,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    deterministic: bool = False,
+    gen: Optional[torch.Generator] = None,
+    rng_state: Optional[torch.Tensor] = None,
+) -> List[torch.Tensor]:
+    """
+    Backward pass for multi-head attention with variable sequence lengths.
+    Args:
+        dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]
+        softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]
+        cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]
+        cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]
+        dq: Optional gradient tensor for queries, same shape as q
+        dk: Optional gradient tensor for keys, same shape as k
+        dv: Optional gradient tensor for values, same shape as v
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        max_seqlen_q: Maximum sequence length for queries
+        max_seqlen_k: Maximum sequence length for keys
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        zero_tensors: Whether to zero tensors before computation
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        deterministic: Whether to use deterministic algorithms
+        gen: Optional random number generator
+        rng_state: Optional RNG state from forward pass
+    Returns:
+        List of tensors: [dq, dk, dv]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    return flash_attn_ops.varlen_bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        p_dropout,
+        softmax_scale,
+        zero_tensors,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        gen,
+        rng_state,
+    )
+def fwd_kvcache(
+    q: torch.Tensor,
+    kcache: torch.Tensor,
+    vcache: torch.Tensor,
+    k: Optional[torch.Tensor] = None,
+    v: Optional[torch.Tensor] = None,
+    seqlens_k: Optional[torch.Tensor] = None,
+    rotary_cos: Optional[torch.Tensor] = None,
+    rotary_sin: Optional[torch.Tensor] = None,
+    cache_batch_idx: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    block_table: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    is_rotary_interleaved: bool = False,
+    num_splits: int = 1,
+) -> List[torch.Tensor]:
+    """
+    Forward pass for multi-head attention with KV cache.
+    Args:
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]
+        v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]
+        seqlens_k: Optional sequence lengths for keys of shape [batch_size]
+        rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2]
+        rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2]
+        cache_batch_idx: Optional indices to index into the KV cache
+        leftpad_k: Optional left padding for keys of shape [batch_size]
+        block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        out: Optional output tensor, same shape as q
+        softmax_scale: Scale factor for softmax
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        is_rotary_interleaved: Whether rotary embeddings are interleaved
+        num_splits: Number of splits for computation
+    Returns:
+        List of tensors: [output, softmax_lse]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    return flash_attn_ops.fwd_kvcache(
+        q,
+        kcache,
+        vcache,
+        k,
+        v,
+        seqlens_k,
+        rotary_cos,
+        rotary_sin,
+        cache_batch_idx,
+        leftpad_k,
+        block_table,
+        alibi_slopes,
+        out,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        is_rotary_interleaved,
+        num_splits,
+    )

build/torch-stable-abi210-cu130-aarch64-linux/_flash_attn2_cuda_460047c.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:5d3349b4e099095c41123643df2cbff49c7065bf03fb5ae0d88f598cd721e413
+size 1005585784

build/torch-stable-abi210-cu130-aarch64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _flash_attn2_cuda_460047c
+ops = torch.ops._flash_attn2_cuda_460047c
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn2_cuda_460047c::{op_name}"

build/torch-stable-abi210-cu130-aarch64-linux/flash_attn_interface.py ADDED Viewed

	@@ -0,0 +1,1708 @@

+# Copyright (c) 2023, Tri Dao.
+from typing import Optional, Sequence, Tuple, Union
+import torch
+import torch.nn as nn
+import os
+# # isort: off
+# # We need to import the CUDA kernels after importing torch
+# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE"
+# if USE_TRITON_ROCM:
+#     from .flash_attn_triton_amd import interface_fa as flash_attn
+# else:
+#     import flash_attn_2_cuda as flash_attn
+from ._ops import ops as flash_attn
+from ._ops import add_op_namespace_prefix
+# # isort: on
+# Torch >= 2.13 can pass a Generator through the stable ABI. Older versions reject the argument
+# before the kernel runs, so there the CUDA op has to draw its Philox seed and offset itself.
+_STABLE_ABI_GENERATOR = torch.__version__ >= "2.13"
+def maybe_contiguous(x):
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+def _get_device():
+    if torch.xpu.is_available():
+        return "xpu"
+    elif torch.cuda.is_available():
+        return "cuda"
+    else:
+        return "cpu"
+def _get_block_size_n(device, head_dim, is_dropout, is_causal):
+    assert head_dim <= 256
+    if device.type == "xpu":
+        if head_dim <= 96:
+            return 64
+        elif head_dim <= 128:
+            return 32
+        elif head_dim <= 256:
+            return 64
+        else:
+            return 32
+    # This should match the block sizes in the CUDA kernel
+    major, minor = torch.cuda.get_device_capability(device)
+    is_sm8x = (
+        major == 8 and minor > 0
+    )  # Only include sm86 and sm89, exclude sm80 (A100)
+    is_sm80 = major == 8 and minor == 0
+    is_sm90 = major == 9 and minor == 0
+    if head_dim <= 32:
+        return 128
+    if head_dim <= 64:
+        return 128 if not is_dropout else 64
+    elif head_dim <= 96:
+        return 64
+    elif head_dim <= 128:
+        if is_sm8x:
+            return 64 if (not is_dropout and is_causal) else 32
+        else:
+            return 64 if not is_dropout else 32
+    elif head_dim <= 192:
+        return 64
+    elif head_dim <= 224:
+        return 64
+    elif head_dim <= 256:
+        return 64
+def round_multiple(x, m):
+    return (x + m - 1) // m * m
+def _philox_generator(
+    q: torch.Tensor, dropout_p: float, generator: Optional[torch.Generator] = None
+) -> Optional[torch.Generator]:
+    """The generator to hand to `fwd` / `varlen_fwd`.
+    With dropout on CUDA this is the default generator unless one is given, so the kernel reserves
+    its seed and offset exactly like the non-stable build (see set_params_philox in flash_api.cpp).
+    """
+    if not q.is_cuda:
+        return generator
+    if dropout_p <= 0:
+        return None
+    if not _STABLE_ABI_GENERATOR:
+        if generator is not None:
+            raise ValueError(
+                "flash_attn2: passing a generator to the CUDA kernel requires torch >= 2.13; "
+                "seed the default generator with torch.manual_seed instead"
+            )
+        return None
+    if generator is None:
+        generator = torch.cuda.default_generators[q.device.index]
+    return generator
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_forward"),
+    mutates_args=(),
+    device_types=_get_device(),
+    # Tells torch.compile the forward draws random numbers, so it is never deduplicated or
+    # recomputed (e.g. under activation checkpointing).
+    tags=(torch.Tag.nondeterministic_seeded,),
+)
+def _flash_attn_forward(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    return_softmax: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    gen = _philox_generator(q, dropout_p)
+    out, softmax_lse, S_dmask, rng_state = flash_attn.fwd(
+        q,
+        k,
+        v,
+        None,
+        alibi_slopes,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out, softmax_lse, S_dmask, rng_state
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward"))
+def _flash_attn_forward_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    return_softmax: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    batch_size, seqlen_q, num_heads, head_size = q.shape
+    seqlen_k = k.shape[1]
+    out = torch.empty_like(q)
+    softmax_lse = torch.empty(
+        (batch_size, num_heads, seqlen_q),
+        dtype=torch.float32,
+        device=q.device,
+        layout=q.layout,
+    )
+    p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout)
+    if return_softmax:
+        p = torch.empty(
+            (
+                batch_size,
+                num_heads,
+                round_multiple(seqlen_q, 128),
+                round_multiple(seqlen_k, 128),
+            ),
+            dtype=q.dtype,
+            device=q.device,
+            layout=q.layout,
+        )
+    rng_state = torch.empty((2,), dtype=torch.int64, device=q.device)
+    return out, softmax_lse, p, rng_state
+_wrapped_flash_attn_forward = flash_attn._flash_attn_forward  # upstream alias, see CARD.md
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_varlen_forward"),
+    mutates_args=(),
+    device_types=_get_device(),
+    tags=(torch.Tag.nondeterministic_seeded,),
+)
+def _flash_attn_varlen_forward(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    return_softmax: bool = False,
+    block_table: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    gen = _philox_generator(q, dropout_p)
+    out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd(
+        q,
+        k,
+        v,
+        None,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_k,
+        leftpad_k,
+        block_table,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        zero_tensors,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    # if out.isnan().any() or softmax_lse.isnan().any():
+    #     breakpoint()
+    return out, softmax_lse, S_dmask, rng_state
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_forward"))
+def _flash_attn_varlen_forward_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    return_softmax: bool = False,
+    block_table: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    paged_kv = block_table is not None
+    batch_size = cu_seqlens_q.numel() - 1
+    total_q, num_heads, _ = q.shape
+    out = torch.empty_like(q)
+    softmax_lse = torch.empty(
+        (num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout
+    )
+    p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout)
+    seqlen_q_rounded = round_multiple(max_seqlen_q, 128)
+    seqlen_k_rounded = round_multiple(max_seqlen_k, 128)
+    if return_softmax:
+        p = torch.empty(
+            (batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded),
+            dtype=q.dtype,
+            device=q.device,
+            layout=q.layout,
+        )
+    rng_state = torch.empty((2,), dtype=torch.int64, device=q.device)
+    return out, softmax_lse, p, rng_state
+_wrapped_flash_attn_varlen_forward = flash_attn._flash_attn_varlen_forward  # upstream alias, see CARD.md
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_backward"),
+    mutates_args=("dq", "dk", "dv"),
+    device_types=_get_device(),
+)
+def _flash_attn_backward(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    # dq, dk, dv are allocated by us so they should already be contiguous
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    (
+        dq,
+        dk,
+        dv,
+        softmax_d,
+    ) = flash_attn.bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        alibi_slopes,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        None,
+        rng_state,
+    )
+    return softmax_d
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward"))
+def _flash_attn_backward_fake(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    if dq is None:
+        dq = torch.empty_like(q)
+    if dk is None:
+        dk = torch.empty_like(k)
+    if dv is None:
+        dv = torch.empty_like(v)
+    batch_size, seqlen_q, num_heads, _ = q.shape
+    softmax_d = torch.empty(
+        (batch_size, num_heads, round_multiple(seqlen_q, 128)),
+        device=q.device,
+        dtype=torch.float32,
+    )
+    return softmax_d
+_wrapped_flash_attn_backward = flash_attn._flash_attn_backward  # upstream alias, see CARD.md
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_varlen_backward"),
+    mutates_args=("dq", "dk", "dv"),
+    device_types=_get_device(),
+)
+def _flash_attn_varlen_backward(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> torch.Tensor:
+    # dq, dk, dv are allocated by us so they should already be contiguous
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    (
+        dq,
+        dk,
+        dv,
+        softmax_d,
+    ) = flash_attn.varlen_bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        zero_tensors,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        None,
+        rng_state,
+    )
+    # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any():
+    #     breakpoint()
+    return softmax_d
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_backward"))
+def _flash_attn_varlen_backward_fake(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> torch.Tensor:
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    batch_size = cu_seqlens_q.numel() - 1
+    total_q, num_heads, _ = q.shape
+    if dq is None:
+        dq = torch.empty_like(q)
+    if dk is None:
+        dk = torch.empty_like(k)
+    if dv is None:
+        dv = torch.empty_like(v)
+    softmax_d = torch.empty(
+        (num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32
+    )
+    return softmax_d
+_wrapped_flash_attn_varlen_backward = flash_attn._flash_attn_varlen_backward  # upstream alias, see CARD.md
+class FlashAttnQKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        qkv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and qkv.requires_grad
+        if softmax_scale is None:
+            softmax_scale = qkv.shape[-1] ** (-0.5)
+        q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach()
+        head_size_og = q.size(3)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+        )
+        if is_grad:
+            ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
+        qkv_shape = q.shape[:-2] + (3, *q.shape[-2:])
+        dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device)
+        head_size_og = dout.size(3)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dqkv[:, :, 0],
+            dqkv[:, :, 1],
+            dqkv[:, :, 2],
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dqkv = dqkv[..., : dout.shape[-1]]  # We could have padded the head dimension
+        return dqkv, None, None, None, None, None, None, None, None, None
+class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        qkv,
+        cu_seqlens,
+        max_seqlen,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and qkv.requires_grad
+        if softmax_scale is None:
+            softmax_scale = qkv.shape[-1] ** (-0.5)
+        q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach()
+        head_size_og = q.size(2)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
+            q,
+            k,
+            v,
+            cu_seqlens,
+            cu_seqlens,
+            max_seqlen,
+            max_seqlen,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+            block_table=None,
+        )
+        if is_grad:
+            ctx.save_for_backward(
+                q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state
+            )
+            ctx.dropout_p = dropout_p
+            ctx.max_seqlen = max_seqlen
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors
+        qkv_shape = q.shape[:-2] + (3, *q.shape[-2:])
+        dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device)
+        head_size_og = dout.size(2)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_varlen_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dqkv[:, 0],
+            dqkv[:, 1],
+            dqkv[:, 2],
+            cu_seqlens,
+            cu_seqlens,
+            ctx.max_seqlen,
+            ctx.max_seqlen,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dqkv = dqkv[..., : dout.shape[-1]]  # We could have padded the head dimension
+        return dqkv, None, None, None, None, None, None, None, None, None, None, None
+class FlashAttnKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        kv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach()
+        head_size_og = q.size(3)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+        )
+        if is_grad:
+            ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
+        dq = torch.empty_like(q)
+        kv_shape = k.shape[:-2] + (2, *k.shape[-2:])
+        dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device)
+        head_size_og = dout.size(3)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dkv[:, :, 0],
+            dkv[:, :, 1],
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dkv = dkv[..., : dout.shape[-1]]
+        return dq, dkv, None, None, None, None, None, None, None, None, None
+class FlashAttnVarlenKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        kv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        k, v = kv[:, 0].detach(), kv[:, 1].detach()
+        head_size_og = q.size(2)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+            block_table=None,
+        )
+        if is_grad:
+            ctx.save_for_backward(
+                q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state
+            )
+            ctx.dropout_p = dropout_p
+            ctx.max_seqlen_q = max_seqlen_q
+            ctx.max_seqlen_k = max_seqlen_k
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = (
+            ctx.saved_tensors
+        )
+        dq = torch.empty_like(q)
+        kv_shape = k.shape[:-2] + (2, *k.shape[-2:])
+        dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device)
+        head_size_og = dout.size(2)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_varlen_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dkv[:, 0],
+            dkv[:, 1],
+            cu_seqlens_q,
+            cu_seqlens_k,
+            ctx.max_seqlen_q,
+            ctx.max_seqlen_k,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dkv = dkv[..., : dout.shape[-1]]
+        return (
+            dq,
+            dkv,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+class FlashAttnFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        head_size_og = q.size(3)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+        )
+        if is_grad:
+            ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
+        dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+        head_size_og = dout.size(3)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dk,
+            dv,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dk = dk[..., : dout.shape[-1]]
+        dv = dv[..., : dout.shape[-1]]
+        return dq, dk, dv, None, None, None, None, None, None, None, None, None
+class FlashAttnVarlenFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        block_table,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        head_size_og = q.size(2)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+            block_table=block_table,
+        )
+        if is_grad:
+            ctx.save_for_backward(
+                q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state
+            )
+            ctx.dropout_p = dropout_p
+            ctx.max_seqlen_q = max_seqlen_q
+            ctx.max_seqlen_k = max_seqlen_k
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = (
+            ctx.saved_tensors
+        )
+        dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+        head_size_og = dout.size(2)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_varlen_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dk,
+            dv,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            ctx.max_seqlen_q,
+            ctx.max_seqlen_k,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dk = dk[..., : dout.shape[-1]]
+        dv = dv[..., : dout.shape[-1]]
+        return (
+            dq,
+            dk,
+            dv,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+def flash_attn_qkvpacked_func(
+    qkv,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # <=0.0 means deactivate
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If Q, K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of Q, K, V.
+    For multi-query and grouped-query attention (MQA/GQA), please see
+    flash_attn_kvpacked_func and flash_attn_func.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.
+    Arguments:
+        qkv: (batch_size, seqlen, 3, nheads, headdim)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to
+            the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnQKVPackedFunc.apply(
+        qkv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_kvpacked_func(
+    q,
+    kv,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of K, V.
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        kv: (batch_size, seqlen, 2, nheads_k, headdim)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnKVPackedFunc.apply(
+        q,
+        kv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_func(
+    q,
+    k,
+    v,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        k: (batch_size, seqlen, nheads_k, headdim)
+        v: (batch_size, seqlen, nheads_k, headdim)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnFunc.apply(
+        q,
+        k,
+        v,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_varlen_qkvpacked_func(
+    qkv,
+    cu_seqlens,
+    max_seqlen,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If Q, K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_varlen_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of Q, K, V.
+    For multi-query and grouped-query attention (MQA/GQA), please see
+    flash_attn_varlen_kvpacked_func and flash_attn_varlen_func.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.
+    Arguments:
+        qkv: (total, 3, nheads, headdim), where total = total number of tokens in the batch.
+        cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into qkv.
+        max_seqlen: int. Maximum sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnVarlenQKVPackedFunc.apply(
+        qkv,
+        cu_seqlens,
+        max_seqlen,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_varlen_kvpacked_func(
+    q,
+    kv,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of K, V.
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.
+        kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into q.
+        cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into kv.
+        max_seqlen_q: int. Maximum query sequence length in the batch.
+        max_seqlen_k: int. Maximum key sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnVarlenKVPackedFunc.apply(
+        q,
+        kv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_varlen_func(
+    q,
+    k,
+    v,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+    block_table=None,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.
+        k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into q.
+        cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into kv.
+        max_seqlen_q: int. Maximum query sequence length in the batch.
+        max_seqlen_k: int. Maximum key sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnVarlenFunc.apply(
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        block_table,
+        False if q.device.type == "cpu" else torch.is_grad_enabled(),
+    )
+def flash_attn_with_kvcache(
+    q,
+    k_cache,
+    v_cache,
+    k=None,
+    v=None,
+    rotary_cos=None,
+    rotary_sin=None,
+    cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None,
+    cache_batch_idx: Optional[torch.Tensor] = None,
+    cache_leftpad: Optional[torch.Tensor] = None,
+    block_table: Optional[torch.Tensor] = None,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    rotary_interleaved=True,
+    alibi_slopes=None,
+    num_splits=0,
+    return_softmax_lse=False,
+):
+    """
+    If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from
+    k and v. This is useful for incremental decoding: you can pass in the cached keys/values from
+    the previous step, and update them with the new keys/values from the current step, and do
+    attention with the updated cache, all in 1 kernel.
+    If you pass in k / v, you must make sure that the cache is large enough to hold the new values.
+    For example, the KV cache could be pre-allocated with the max sequence length, and you can use
+    cache_seqlens to keep track of the current sequence lengths of each sequence in the batch.
+    Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be
+    rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
+    If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos
+    and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
+    If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at
+    indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens).
+    See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function.
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Note: Does not support backward pass.
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,
+            or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)
+            page_block_size must be a multiple of 256.
+        v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,
+            or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)
+        k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate
+            k with k_cache, starting at the indices specified by cache_seqlens.
+        v [optional]: (batch_size, seqlen_new, nheads_k, headdim). Similar to k.
+        rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding
+            to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16.
+        rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos.
+        cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the
+            KV cache.
+        cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache.
+            If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1].
+            If the indices are not distinct, and k and v are provided, the values updated in the cache
+                 might come from any of the duplicate indices.
+        cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0.
+        block_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in.
+            If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False,
+            rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1
+            (i.e. GPT-NeoX style).
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        num_splits: int. If > 1, split the key/value into this many chunks along the sequence.
+           If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic
+           to automatically determine the number of splits.
+           Don't change this unless you know what you are doing.
+        return_softmax_lse: bool. Whether to return the logsumexp of the attention scores.
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+    """
+    assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension"
+    assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension"
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** (-0.5)
+    if cache_seqlens is not None and isinstance(cache_seqlens, int):
+        cache_seqlens = torch.full(
+            (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device
+        )
+        cache_seqlens = maybe_contiguous(cache_seqlens)
+    cache_batch_idx = maybe_contiguous(cache_batch_idx)
+    block_table = maybe_contiguous(block_table)
+    out, softmax_lse = flash_attn.fwd_kvcache(
+        q,
+        k_cache,
+        v_cache,
+        k,
+        v,
+        cache_seqlens,
+        rotary_cos,
+        rotary_sin,
+        cache_batch_idx,
+        cache_leftpad,
+        block_table,
+        alibi_slopes,
+        None,
+        softmax_scale,
+        causal,
+        window_size[0],
+        window_size[1],
+        softcap,
+        rotary_interleaved,
+        num_splits,
+    )
+    return (out, softmax_lse) if return_softmax_lse else out

build/torch-stable-abi210-cu130-aarch64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,39 @@

+{
+  "name": "flash-attn2",
+  "id": "_flash_attn2_cuda_460047c",
+  "version": 4,
+  "kernels-minver": "0.14.0",
+  "license": "BSD-3-Clause",
+  "python-depends": [],
+  "kernel-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Nz8hqrusjoUHQqszpxfF2USMwUXn1VfpvcmtAAvutSM=",
+      "_flash_attn2_cuda_460047c.abi3.so": "XTNJtOCZCVxBEjZD3yy/9JxwZb8D+1rg2I9ZjNch5BM=",
+      "_ops.py": "DV4Eo4+s3aaNmDsfXoZiYKNLnRUJqlsdDGlMsdKFgk4=",
+      "flash_attn_interface.py": "6UYUhfUax4+yOUygyS7K8dVjrIjU5GQr1AJ2gSp8rPo=",
+      "symbols.json": "uvC+87tHUhtHv7L/Nl53T/Ms7ceX317/jwQBoqHY14A="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.18.0-dev0",
+      "commit": "d054906a13710b811b96697b4635f03020842733",
+      "dirty": false
+    },
+    "kernel": {
+      "commit": "460047c291fc7f0b3dbd4239b46bd5e1281b1806",
+      "dirty": false
+    }
+  }
+}

build/torch-stable-abi210-cu130-aarch64-linux/metadata.json.sigstore ADDED Viewed

	@@ -0,0 +1 @@

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrOgAwIBAgIUV5dC/QsMY1gdXh+pxvVOTFf2nxQwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDI0MDAwNzUwWhcNMjYxMDI0MDAxNzUwWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEpYnhGECr0UrvaU7FKpQzhY0NMJ7+WKOgdTFvK15HSzJ6bI193348JWbgHhcnDqVZplaBy/PdUwCjoomFQVl6UaOCBdIwggXOMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBCg0NjAwNDdjMjkxZmM3ZjBiM2RiZDQyMzliNDZiZDVlMTI4MWIxODA2MBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDCg0NjAwNDdjMjkxZmM3ZjBiM2RiZDQyMzliNDZiZDVlMTI4MWIxODA2MBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDCg0NjAwNDdjMjkxZmM3ZjBiM2RiZDQyMzliNDZiZDVlMTI4MWIxODA2MB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoNDYwMDQ3YzI5MWZjN2YwYjNkYmQ0MjM5YjQ2YmQ1ZTEyODFiMTgwNjAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTk4MjY0MDUxNS9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGLBgorBgEEAdZ5AgQCBH0EewB5AHcA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhazw2HgAABAMASDBGAiEA/sdxrIoHwWLtLYH/zSmq7ZuJVtwcZhFcfeuCwtq4zaYCIQD9DBjsxqDhWas4nf3qZWRPN+4RkY8nlDO6DNPFLkhEBDAKBggqhkjOPQQDAwNnADBkAjBPzCJZB4fe4+gup4iAsdCMJDS2r6ytH3dy6FMauWIXJUdzHh4tIY1jmNIy0wOMsDkCMA7GDNNIS0J/4woDrEUlbZmH035vUjLltFZuK+IQsbv7a91yAUhQxRM3Zk/hUId3ZA=="},"tlogEntries":[{"logIndex":"3430231988","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792800470","inclusionPromise":{"signedEntryTimestamp":"MEYCIQCmHkPxEo6UA56CtGAvcd7tuGSq+BTh09PAdp32SFBZLgIhAPToHBO342Kcu8F0o+BGUR+W22R/ILnHLl73H3vqP8Lj"},"inclusionProof":{"logIndex":"3308327726","rootHash":"9j+WvjjLBQ6Dg17XvdJ7GaT9NstggA0SjQdvQa4SJls=","treeSize":"3308328149","hashes":["NMP7eOc8EnoNjTnvyncEDeeDSn30GrZLCo/z/YmE8Lw=","ICRVwZIV8m9NxN75D8+IxvBmuVn1gqs5Xh10I4RFKY0=","WBxEYgVoTGEhcWfbUQduXtRHERE8IK0oWSPwDHRLT/U=","MnnalYww5Z3KevmaBMHkF5DOiCuZDATeV5EhjKf50LM=","Tc0I1OuGp5qzxaeLQEWZyOeCB8SHUPvP0eNG/K38o2c=","YwDNh7SwtIh/zqCVP1jL0CyyakPPls+pxGBrL4UWnKU=","qg3u2EeN52cnvOm+KIKO7RMKmtx8rYCMwXusD5VORnc=","+9D72OvJMqE7pujmt6xkkyNsnwEo55yUuDk9sn3o+eY=","XjdWy49UvJFoNoQxicEvigURdRnav3eX5vfaEMQxdR0=","C4d/JSc723ScaFwqOj7Rp4UyeZlWzqzRqTehzamGn7Y=","hM/G+D/5k6TXc6r9PyzYqE/YqDRL1txJXgyl5W+peYg=","gdS39fsn5MOVlKluUgGY2sdKF3H5oauxkHMa7DSSNcg=","8x0N/oAmvpY+6DQUGTU+i6a1Cuqphro2WbV5WZ/RHL8=","DMchGrSAeKV9qHo75XJV3UFZ/CHF8zofPRu+NgVRwA4=","V6dXjeDcwYxBG/9gVSBM8rYdPrAqGpZsf1e7v6l67GA=","iDfImp65w8dirRfYGNgFA34JTUYDWbcJwNBBU/Ok+X4=","YYCr61A8azar3dfO8Ct+pri7LqDxkUlUkfs7fiqYErU=","1DaDKzrCNNBhmSdY5vFjz+NpAmxmYfYEgBCdPLrMP1A=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3308328149\n9j+WvjjLBQ6Dg17XvdJ7GaT9NstggA0SjQdvQa4SJls=\n\n— rekor.sigstore.dev wNI9ajBEAiBZmvNl49SFK82hdRPtmAHH84wBRtvOKW++tQp9p62duwIga+rvFXCpKRulXzsOALA3vQv+xB+oJAiUe5ehKu+mpn4=\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI3NjFhZjcxMTBkYjY5OThiODVhZGViYWYwMjdhZTA5MDI4ZDA5ZmNmNWUzYWI5MzMyM2VmMDU3NjA3MzFjODEzIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJSGZsUUlkZlh6dHVFM2V5VzFiQUdidkZGR1RsajZ2L3VsYlFCN3JHRVhvT0FpQWswRUs4aDc0NzluVUI3KzVvMVhleCtraTFmMUtCNlJhaUJ2TXBJOVVSdXc9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KUFowRjNTVUpCWjBsVlZqVmtReTlSYzAxWk1XZGtXR2dyY0hoMlZrOVVSbVl5Ym5oUmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFU1RCTlJFRjNUbnBWZDFkb1kwNU5hbGw0VFVSSk1FMUVRWGhPZWxWM1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ3V1c1b1IwVkRjakJWY25aaFZUZEdTM0JSZW1oWk1FNU5TamNyVjB0UFoyUlVSbllLU3pFMVNGTjZTalppU1RFNU16TTBPRXBYWW1kSWFHTnVSSEZXV25Cc1lVSjVMMUJrVlhkRGFtOXZiVVpSVm13MlZXRlBRMEprU1hkbloxaFBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmN3VG1wQmR3cE9SR1JxVFdwcmVGcHRUVE5hYWtKcFRUSlNhVnBFVVhsTmVteHBUa1JhYVZwRVZteE5WRWswVFZkSmVFOUVRVEpOUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEWnpCT2FrRjNUa1JrYWsxcWEzaGFiVTB6Q2xwcVFtbE5NbEpwV2tSUmVVMTZiR2xPUkZwcFdrUldiRTFVU1RSTlYwbDRUMFJCTWsxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEWnpCT2FrRjNDazVFWkdwTmFtdDRXbTFOTTFwcVFtbE5NbEpwV2tSUmVVMTZiR2xPUkZwcFdrUldiRTFVU1RSTlYwbDRUMFJCTWsxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlPUkZsM1RVUlJNMWw2U1RWTlYxcHFUakpaZDFscVRtdFpiVkV3VFdwTk5WbHFVVEpaYlZFeFdsUkZlVTlFUm1rS1RWUm5kMDVxUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUnJORTFxV1RCTlJGVjRUbE01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTUNrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NEQkZaWGRDTlVGSVkwRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvWVhwM01raG5RVUZDUVUxQlUwUkNSMEZwUlVFdmMyUjRja2x2U0hkWFRIUk1XVWd2ZWxOdGNUZGFkVXBXZEhkaldtaEdZd3BtWlhWRGQzUnhOSHBoV1VOSlVVUTVSRUpxYzNoeFJHaFhZWE0wYm1ZemNWcFhVbEJPS3pSU2ExazRibXhFVHpaRVRsQkdUR3RvUlVKRVFVdENaMmR4Q21ocmFrOVFVVkZFUVhkT2JrRkVRbXRCYWtKUWVrTktXa0kwWm1VMEsyZDFjRFJwUVhOa1EwMUtSRk15Y2paNWRFZ3paSGsyUmsxaGRWZEpXRXBWWkhvS1NHZzBkRWxaTVdwdFRrbDVNSGRQVFhORWEwTk5RVGRIUkU1T1NWTXdTaTgwZDI5RWNrVlZiR0phYlVnd016VjJWV3BNYkhSR1duVkxLMGxSYzJKMk53cGhPVEY1UVZWb1VYaFNUVE5hYXk5b1ZVbGtNMXBCUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg2QDIXWO9VN8DxxfI85ni4MORa8cBbgJH6V5ccpOXqAICFQCVAWSZuF6tN+YWnLggZoOKD7AHMRgPMjAyNjEwMjQwMDA3NTBaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAyNDAwMDc1MFowLwYJKoZIhvcNAQkEMSIEILc7Pdnq0+D+Wv63pX98gN3xy67OZdkEQ65gYLOVewlGMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQDZQTHmRv2fKD0ezY8hk7lM1sHXqDGsSEzg0Bv4I+zLjAoMOxk1sXFqPJzlAZQCBMkCMQDwXXtXkhgqlcXEdGdUQ07XYZS90rGgRNiBS0VauRqvEMKCSPehugeIoH5zoaBHYa8="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"dhr3EQ22mYuFreuvAnrgkCjQn89eOrkzI+8FdgcxyBM="},"signature":"MEQCIHflQIdfXztuE3eyW1bAGbvFFGTlj6v/ulbQB7rGEXoOAiAk0EK8h7479nUB7+5o1Xex+ki1f1KB6RaiBvMpI9URuw=="}}

build/torch-stable-abi210-cu130-aarch64-linux/symbols.json ADDED Viewed

	@@ -0,0 +1,1273 @@

+{
+  "schema_version": 1,
+  "module": "_flash_attn2_cuda_460047c",
+  "functions": [
+    {
+      "name": "flash_attn_func",
+      "qualname": "flash_attn_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (batch_size, seqlen, nheads, headdim)\n    k: (batch_size, seqlen, nheads_k, headdim)\n    v: (batch_size, seqlen, nheads_k, headdim)\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_kvpacked_func",
+      "qualname": "flash_attn_kvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of K, V.\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (batch_size, seqlen, nheads, headdim)\n    kv: (batch_size, seqlen, 2, nheads_k, headdim)\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "kv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_qkvpacked_func",
+      "qualname": "flash_attn_qkvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf Q, K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of Q, K, V.\nFor multi-query and grouped-query attention (MQA/GQA), please see\nflash_attn_kvpacked_func and flash_attn_func.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.\n\nArguments:\n    qkv: (batch_size, seqlen, 3, nheads, headdim)\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to\n        the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "qkv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_varlen_func",
+      "qualname": "flash_attn_varlen_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.\n    k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.\n    v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.\n    cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into q.\n    cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into kv.\n    max_seqlen_q: int. Maximum query sequence length in the batch.\n    max_seqlen_k: int. Maximum key sequence length in the batch.\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (total, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_varlen_kvpacked_func",
+      "qualname": "flash_attn_varlen_kvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of K, V.\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.\n    kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch.\n    cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into q.\n    cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into kv.\n    max_seqlen_q: int. Maximum query sequence length in the batch.\n    max_seqlen_k: int. Maximum key sequence length in the batch.\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (total, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "kv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_varlen_qkvpacked_func",
+      "qualname": "flash_attn_varlen_qkvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf Q, K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_varlen_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of Q, K, V.\nFor multi-query and grouped-query attention (MQA/GQA), please see\nflash_attn_varlen_kvpacked_func and flash_attn_varlen_func.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.\n\nArguments:\n    qkv: (total, 3, nheads, headdim), where total = total number of tokens in the batch.\n    cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into qkv.\n    max_seqlen: int. Maximum sequence length in the batch.\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (total, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "qkv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_with_kvcache",
+      "qualname": "flash_attn_with_kvcache",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from\nk and v. This is useful for incremental decoding: you can pass in the cached keys/values from\nthe previous step, and update them with the new keys/values from the current step, and do\nattention with the updated cache, all in 1 kernel.\n\nIf you pass in k / v, you must make sure that the cache is large enough to hold the new values.\nFor example, the KV cache could be pre-allocated with the max sequence length, and you can use\ncache_seqlens to keep track of the current sequence lengths of each sequence in the batch.\n\nAlso apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be\nrotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.\nIf causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos\nand rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.\nIf not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at\nindices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens).\n\nSee tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function.\n\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nNote: Does not support backward pass.\n\nArguments:\n    q: (batch_size, seqlen, nheads, headdim)\n    k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,\n        or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)\n        page_block_size must be a multiple of 256.\n    v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,\n        or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)\n    k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate\n        k with k_cache, starting at the indices specified by cache_seqlens.\n    v [optional]: (batch_size, seqlen_new, nheads_k, headdim). Similar to k.\n    rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding\n        to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16.\n    rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos.\n    cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the\n        KV cache.\n    cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache.\n        If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1].\n        If the indices are not distinct, and k and v are provided, the values updated in the cache\n             might come from any of the duplicate indices.\n    cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0.\n    block_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in.\n        If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False,\n        rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1\n        (i.e. GPT-NeoX style).\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    num_splits: int. If > 1, split the key/value into this many chunks along the sequence.\n       If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic\n       to automatically determine the number of splits.\n       Don't change this unless you know what you are doing.\n    return_softmax_lse: bool. Whether to return the logsumexp of the attention scores.\n\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k_cache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "v_cache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "rotary_cos",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "rotary_sin",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "cache_seqlens",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "int | torch.Tensor | None"
+          },
+          {
+            "name": "cache_batch_idx",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "cache_leftpad",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "rotary_interleaved",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "True",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "num_splits",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": null
+          },
+          {
+            "name": "return_softmax_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "fwd",
+      "qualname": "fwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Forward pass for multi-head attention.\n\nArgs:\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    out: Optional output tensor, same shape as q\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    return_softmax: Whether to return softmax weights\n    gen: Optional random number generator\n\nReturns:\n    List of tensors: [output, softmax_lse, (softmax if return_softmax)]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "return_softmax",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "varlen_fwd",
+      "qualname": "varlen_fwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Forward pass for multi-head attention with variable sequence lengths.\n\nArgs:\n    q: Query tensor of shape [total_q, num_heads, head_size]\n    k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]\n    cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]\n    out: Optional output tensor of shape [total_q, num_heads, head_size]\n    seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size]\n    leftpad_k: Optional left padding for keys of shape [batch_size]\n    block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    max_seqlen_q: Maximum sequence length for queries\n    max_seqlen_k: Maximum sequence length for keys\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    zero_tensors: Whether to zero tensors before computation\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    return_softmax: Whether to return softmax weights\n    gen: Optional random number generator\n\nReturns:\n    List of tensors: [output, softmax_lse, (softmax if return_softmax)]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "seqused_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "leftpad_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "zero_tensors",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "return_softmax",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "bwd",
+      "qualname": "bwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Backward pass for multi-head attention.\n\nArgs:\n    dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]\n    softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]\n    dq: Optional gradient tensor for queries, same shape as q\n    dk: Optional gradient tensor for keys, same shape as k\n    dv: Optional gradient tensor for values, same shape as v\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    deterministic: Whether to use deterministic algorithms\n    gen: Optional random number generator\n    rng_state: Optional RNG state from forward pass\n\nReturns:\n    List of tensors: [dq, dk, dv]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "dout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "softmax_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "dq",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dk",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          },
+          {
+            "name": "rng_state",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "varlen_bwd",
+      "qualname": "varlen_bwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Backward pass for multi-head attention with variable sequence lengths.\n\nArgs:\n    dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]\n    softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]\n    cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]\n    cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]\n    dq: Optional gradient tensor for queries, same shape as q\n    dk: Optional gradient tensor for keys, same shape as k\n    dv: Optional gradient tensor for values, same shape as v\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    max_seqlen_q: Maximum sequence length for queries\n    max_seqlen_k: Maximum sequence length for keys\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    zero_tensors: Whether to zero tensors before computation\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    deterministic: Whether to use deterministic algorithms\n    gen: Optional random number generator\n    rng_state: Optional RNG state from forward pass\n\nReturns:\n    List of tensors: [dq, dk, dv]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "dout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "softmax_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "dq",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dk",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "zero_tensors",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          },
+          {
+            "name": "rng_state",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "fwd_kvcache",
+      "qualname": "fwd_kvcache",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Forward pass for multi-head attention with KV cache.\n\nArgs:\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]\n    v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]\n    seqlens_k: Optional sequence lengths for keys of shape [batch_size]\n    rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2]\n    rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2]\n    cache_batch_idx: Optional indices to index into the KV cache\n    leftpad_k: Optional left padding for keys of shape [batch_size]\n    block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    out: Optional output tensor, same shape as q\n    softmax_scale: Scale factor for softmax\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    is_rotary_interleaved: Whether rotary embeddings are interleaved\n    num_splits: Number of splits for computation\n\nReturns:\n    List of tensors: [output, softmax_lse]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "kcache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "vcache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "rotary_cos",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "rotary_sin",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "cache_batch_idx",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "leftpad_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "is_rotary_interleaved",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "num_splits",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "1",
+            "annotation": "int"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    }
+  ],
+  "layers": []
+}

build/torch-stable-abi210-cu132-aarch64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,424 @@

+from typing import List, Optional
+import torch
+from . import flash_attn_interface
+from ._ops import ops as flash_attn_ops
+from .flash_attn_interface import (
+    # helper to set up the Philox random number generator for dropout
+    _philox_generator,
+    # reexported for convenience to be used in the high-level API functions
+    flash_attn_func,
+    flash_attn_kvpacked_func,
+    flash_attn_qkvpacked_func,
+    flash_attn_varlen_func,
+    flash_attn_varlen_kvpacked_func,
+    flash_attn_varlen_qkvpacked_func,
+    flash_attn_with_kvcache,
+)
+__all__ = [
+    # High-level API (with autograd support)
+    "flash_attn_func",
+    "flash_attn_kvpacked_func",
+    "flash_attn_qkvpacked_func",
+    "flash_attn_varlen_func",
+    "flash_attn_varlen_kvpacked_func",
+    "flash_attn_varlen_qkvpacked_func",
+    "flash_attn_with_kvcache",
+    # Upstream interface module. Its `_flash_attn_forward`, `_flash_attn_backward`,
+    # `_flash_attn_varlen_forward` and `_flash_attn_varlen_backward` are supported
+    # low-level entry points with upstream's signatures: Torch custom ops without
+    # autograd that return the LSE and are traceable by torch.compile. The autograd
+    # Functions above are built on top of them.
+    "flash_attn_interface",
+    # Low-level ops (no autograd)
+    "fwd",
+    "varlen_fwd",
+    "bwd",
+    "varlen_bwd",
+    "fwd_kvcache",
+]
+def fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    return_softmax: bool = False,
+    gen: Optional[torch.Generator] = None,
+) -> List[torch.Tensor]:
+    """
+    Forward pass for multi-head attention.
+    Args:
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        out: Optional output tensor, same shape as q
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        return_softmax: Whether to return softmax weights
+        gen: Optional random number generator
+    Returns:
+        List of tensors: [output, softmax_lse, (softmax if return_softmax)]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    gen = _philox_generator(q, p_dropout, gen)
+    return flash_attn_ops.fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+def varlen_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    block_table: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    max_seqlen_q: int = 0,
+    max_seqlen_k: int = 0,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    zero_tensors: bool = False,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    return_softmax: bool = False,
+    gen: Optional[torch.Generator] = None,
+) -> List[torch.Tensor]:
+    """
+    Forward pass for multi-head attention with variable sequence lengths.
+    Args:
+        q: Query tensor of shape [total_q, num_heads, head_size]
+        k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]
+        cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]
+        out: Optional output tensor of shape [total_q, num_heads, head_size]
+        seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size]
+        leftpad_k: Optional left padding for keys of shape [batch_size]
+        block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        max_seqlen_q: Maximum sequence length for queries
+        max_seqlen_k: Maximum sequence length for keys
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        zero_tensors: Whether to zero tensors before computation
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        return_softmax: Whether to return softmax weights
+        gen: Optional random number generator
+    Returns:
+        List of tensors: [output, softmax_lse, (softmax if return_softmax)]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    gen = _philox_generator(q, p_dropout, gen)
+    return flash_attn_ops.varlen_fwd(
+        q,
+        k,
+        v,
+        out,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_k,
+        leftpad_k,
+        block_table,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        p_dropout,
+        softmax_scale,
+        zero_tensors,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+def bwd(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor] = None,
+    dk: Optional[torch.Tensor] = None,
+    dv: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    deterministic: bool = False,
+    gen: Optional[torch.Generator] = None,
+    rng_state: Optional[torch.Tensor] = None,
+) -> List[torch.Tensor]:
+    """
+    Backward pass for multi-head attention.
+    Args:
+        dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]
+        softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]
+        dq: Optional gradient tensor for queries, same shape as q
+        dk: Optional gradient tensor for keys, same shape as k
+        dv: Optional gradient tensor for values, same shape as v
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        deterministic: Whether to use deterministic algorithms
+        gen: Optional random number generator
+        rng_state: Optional RNG state from forward pass
+    Returns:
+        List of tensors: [dq, dk, dv]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    return flash_attn_ops.bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        gen,
+        rng_state,
+    )
+def varlen_bwd(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    dq: Optional[torch.Tensor] = None,
+    dk: Optional[torch.Tensor] = None,
+    dv: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    max_seqlen_q: int = 0,
+    max_seqlen_k: int = 0,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    zero_tensors: bool = False,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    deterministic: bool = False,
+    gen: Optional[torch.Generator] = None,
+    rng_state: Optional[torch.Tensor] = None,
+) -> List[torch.Tensor]:
+    """
+    Backward pass for multi-head attention with variable sequence lengths.
+    Args:
+        dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]
+        softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]
+        cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]
+        cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]
+        dq: Optional gradient tensor for queries, same shape as q
+        dk: Optional gradient tensor for keys, same shape as k
+        dv: Optional gradient tensor for values, same shape as v
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        max_seqlen_q: Maximum sequence length for queries
+        max_seqlen_k: Maximum sequence length for keys
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        zero_tensors: Whether to zero tensors before computation
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        deterministic: Whether to use deterministic algorithms
+        gen: Optional random number generator
+        rng_state: Optional RNG state from forward pass
+    Returns:
+        List of tensors: [dq, dk, dv]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    return flash_attn_ops.varlen_bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        p_dropout,
+        softmax_scale,
+        zero_tensors,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        gen,
+        rng_state,
+    )
+def fwd_kvcache(
+    q: torch.Tensor,
+    kcache: torch.Tensor,
+    vcache: torch.Tensor,
+    k: Optional[torch.Tensor] = None,
+    v: Optional[torch.Tensor] = None,
+    seqlens_k: Optional[torch.Tensor] = None,
+    rotary_cos: Optional[torch.Tensor] = None,
+    rotary_sin: Optional[torch.Tensor] = None,
+    cache_batch_idx: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    block_table: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    is_rotary_interleaved: bool = False,
+    num_splits: int = 1,
+) -> List[torch.Tensor]:
+    """
+    Forward pass for multi-head attention with KV cache.
+    Args:
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]
+        v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]
+        seqlens_k: Optional sequence lengths for keys of shape [batch_size]
+        rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2]
+        rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2]
+        cache_batch_idx: Optional indices to index into the KV cache
+        leftpad_k: Optional left padding for keys of shape [batch_size]
+        block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        out: Optional output tensor, same shape as q
+        softmax_scale: Scale factor for softmax
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        is_rotary_interleaved: Whether rotary embeddings are interleaved
+        num_splits: Number of splits for computation
+    Returns:
+        List of tensors: [output, softmax_lse]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    return flash_attn_ops.fwd_kvcache(
+        q,
+        kcache,
+        vcache,
+        k,
+        v,
+        seqlens_k,
+        rotary_cos,
+        rotary_sin,
+        cache_batch_idx,
+        leftpad_k,
+        block_table,
+        alibi_slopes,
+        out,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        is_rotary_interleaved,
+        num_splits,
+    )

build/torch-stable-abi210-cu132-aarch64-linux/_flash_attn2_cuda_460047c.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:a3829f9e740fcc72ea67cd5334a6917943be836f86bd652957fd62f1e3aea3b9
+size 1021117888

build/torch-stable-abi210-cu132-aarch64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _flash_attn2_cuda_460047c
+ops = torch.ops._flash_attn2_cuda_460047c
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn2_cuda_460047c::{op_name}"

build/torch-stable-abi210-cu132-aarch64-linux/flash_attn_interface.py ADDED Viewed

	@@ -0,0 +1,1708 @@

+# Copyright (c) 2023, Tri Dao.
+from typing import Optional, Sequence, Tuple, Union
+import torch
+import torch.nn as nn
+import os
+# # isort: off
+# # We need to import the CUDA kernels after importing torch
+# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE"
+# if USE_TRITON_ROCM:
+#     from .flash_attn_triton_amd import interface_fa as flash_attn
+# else:
+#     import flash_attn_2_cuda as flash_attn
+from ._ops import ops as flash_attn
+from ._ops import add_op_namespace_prefix
+# # isort: on
+# Torch >= 2.13 can pass a Generator through the stable ABI. Older versions reject the argument
+# before the kernel runs, so there the CUDA op has to draw its Philox seed and offset itself.
+_STABLE_ABI_GENERATOR = torch.__version__ >= "2.13"
+def maybe_contiguous(x):
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+def _get_device():
+    if torch.xpu.is_available():
+        return "xpu"
+    elif torch.cuda.is_available():
+        return "cuda"
+    else:
+        return "cpu"
+def _get_block_size_n(device, head_dim, is_dropout, is_causal):
+    assert head_dim <= 256
+    if device.type == "xpu":
+        if head_dim <= 96:
+            return 64
+        elif head_dim <= 128:
+            return 32
+        elif head_dim <= 256:
+            return 64
+        else:
+            return 32
+    # This should match the block sizes in the CUDA kernel
+    major, minor = torch.cuda.get_device_capability(device)
+    is_sm8x = (
+        major == 8 and minor > 0
+    )  # Only include sm86 and sm89, exclude sm80 (A100)
+    is_sm80 = major == 8 and minor == 0
+    is_sm90 = major == 9 and minor == 0
+    if head_dim <= 32:
+        return 128
+    if head_dim <= 64:
+        return 128 if not is_dropout else 64
+    elif head_dim <= 96:
+        return 64
+    elif head_dim <= 128:
+        if is_sm8x:
+            return 64 if (not is_dropout and is_causal) else 32
+        else:
+            return 64 if not is_dropout else 32
+    elif head_dim <= 192:
+        return 64
+    elif head_dim <= 224:
+        return 64
+    elif head_dim <= 256:
+        return 64
+def round_multiple(x, m):
+    return (x + m - 1) // m * m
+def _philox_generator(
+    q: torch.Tensor, dropout_p: float, generator: Optional[torch.Generator] = None
+) -> Optional[torch.Generator]:
+    """The generator to hand to `fwd` / `varlen_fwd`.
+    With dropout on CUDA this is the default generator unless one is given, so the kernel reserves
+    its seed and offset exactly like the non-stable build (see set_params_philox in flash_api.cpp).
+    """
+    if not q.is_cuda:
+        return generator
+    if dropout_p <= 0:
+        return None
+    if not _STABLE_ABI_GENERATOR:
+        if generator is not None:
+            raise ValueError(
+                "flash_attn2: passing a generator to the CUDA kernel requires torch >= 2.13; "
+                "seed the default generator with torch.manual_seed instead"
+            )
+        return None
+    if generator is None:
+        generator = torch.cuda.default_generators[q.device.index]
+    return generator
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_forward"),
+    mutates_args=(),
+    device_types=_get_device(),
+    # Tells torch.compile the forward draws random numbers, so it is never deduplicated or
+    # recomputed (e.g. under activation checkpointing).
+    tags=(torch.Tag.nondeterministic_seeded,),
+)
+def _flash_attn_forward(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    return_softmax: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    gen = _philox_generator(q, dropout_p)
+    out, softmax_lse, S_dmask, rng_state = flash_attn.fwd(
+        q,
+        k,
+        v,
+        None,
+        alibi_slopes,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out, softmax_lse, S_dmask, rng_state
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward"))
+def _flash_attn_forward_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    return_softmax: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    batch_size, seqlen_q, num_heads, head_size = q.shape
+    seqlen_k = k.shape[1]
+    out = torch.empty_like(q)
+    softmax_lse = torch.empty(
+        (batch_size, num_heads, seqlen_q),
+        dtype=torch.float32,
+        device=q.device,
+        layout=q.layout,
+    )
+    p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout)
+    if return_softmax:
+        p = torch.empty(
+            (
+                batch_size,
+                num_heads,
+                round_multiple(seqlen_q, 128),
+                round_multiple(seqlen_k, 128),
+            ),
+            dtype=q.dtype,
+            device=q.device,
+            layout=q.layout,
+        )
+    rng_state = torch.empty((2,), dtype=torch.int64, device=q.device)
+    return out, softmax_lse, p, rng_state
+_wrapped_flash_attn_forward = flash_attn._flash_attn_forward  # upstream alias, see CARD.md
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_varlen_forward"),
+    mutates_args=(),
+    device_types=_get_device(),
+    tags=(torch.Tag.nondeterministic_seeded,),
+)
+def _flash_attn_varlen_forward(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    return_softmax: bool = False,
+    block_table: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    gen = _philox_generator(q, dropout_p)
+    out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd(
+        q,
+        k,
+        v,
+        None,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_k,
+        leftpad_k,
+        block_table,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        zero_tensors,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    # if out.isnan().any() or softmax_lse.isnan().any():
+    #     breakpoint()
+    return out, softmax_lse, S_dmask, rng_state
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_forward"))
+def _flash_attn_varlen_forward_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    return_softmax: bool = False,
+    block_table: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    paged_kv = block_table is not None
+    batch_size = cu_seqlens_q.numel() - 1
+    total_q, num_heads, _ = q.shape
+    out = torch.empty_like(q)
+    softmax_lse = torch.empty(
+        (num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout
+    )
+    p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout)
+    seqlen_q_rounded = round_multiple(max_seqlen_q, 128)
+    seqlen_k_rounded = round_multiple(max_seqlen_k, 128)
+    if return_softmax:
+        p = torch.empty(
+            (batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded),
+            dtype=q.dtype,
+            device=q.device,
+            layout=q.layout,
+        )
+    rng_state = torch.empty((2,), dtype=torch.int64, device=q.device)
+    return out, softmax_lse, p, rng_state
+_wrapped_flash_attn_varlen_forward = flash_attn._flash_attn_varlen_forward  # upstream alias, see CARD.md
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_backward"),
+    mutates_args=("dq", "dk", "dv"),
+    device_types=_get_device(),
+)
+def _flash_attn_backward(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    # dq, dk, dv are allocated by us so they should already be contiguous
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    (
+        dq,
+        dk,
+        dv,
+        softmax_d,
+    ) = flash_attn.bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        alibi_slopes,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        None,
+        rng_state,
+    )
+    return softmax_d
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward"))
+def _flash_attn_backward_fake(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    if dq is None:
+        dq = torch.empty_like(q)
+    if dk is None:
+        dk = torch.empty_like(k)
+    if dv is None:
+        dv = torch.empty_like(v)
+    batch_size, seqlen_q, num_heads, _ = q.shape
+    softmax_d = torch.empty(
+        (batch_size, num_heads, round_multiple(seqlen_q, 128)),
+        device=q.device,
+        dtype=torch.float32,
+    )
+    return softmax_d
+_wrapped_flash_attn_backward = flash_attn._flash_attn_backward  # upstream alias, see CARD.md
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_varlen_backward"),
+    mutates_args=("dq", "dk", "dv"),
+    device_types=_get_device(),
+)
+def _flash_attn_varlen_backward(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> torch.Tensor:
+    # dq, dk, dv are allocated by us so they should already be contiguous
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    (
+        dq,
+        dk,
+        dv,
+        softmax_d,
+    ) = flash_attn.varlen_bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        zero_tensors,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        None,
+        rng_state,
+    )
+    # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any():
+    #     breakpoint()
+    return softmax_d
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_backward"))
+def _flash_attn_varlen_backward_fake(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> torch.Tensor:
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    batch_size = cu_seqlens_q.numel() - 1
+    total_q, num_heads, _ = q.shape
+    if dq is None:
+        dq = torch.empty_like(q)
+    if dk is None:
+        dk = torch.empty_like(k)
+    if dv is None:
+        dv = torch.empty_like(v)
+    softmax_d = torch.empty(
+        (num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32
+    )
+    return softmax_d
+_wrapped_flash_attn_varlen_backward = flash_attn._flash_attn_varlen_backward  # upstream alias, see CARD.md
+class FlashAttnQKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        qkv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and qkv.requires_grad
+        if softmax_scale is None:
+            softmax_scale = qkv.shape[-1] ** (-0.5)
+        q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach()
+        head_size_og = q.size(3)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+        )
+        if is_grad:
+            ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
+        qkv_shape = q.shape[:-2] + (3, *q.shape[-2:])
+        dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device)
+        head_size_og = dout.size(3)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dqkv[:, :, 0],
+            dqkv[:, :, 1],
+            dqkv[:, :, 2],
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dqkv = dqkv[..., : dout.shape[-1]]  # We could have padded the head dimension
+        return dqkv, None, None, None, None, None, None, None, None, None
+class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        qkv,
+        cu_seqlens,
+        max_seqlen,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and qkv.requires_grad
+        if softmax_scale is None:
+            softmax_scale = qkv.shape[-1] ** (-0.5)
+        q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach()
+        head_size_og = q.size(2)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
+            q,
+            k,
+            v,
+            cu_seqlens,
+            cu_seqlens,
+            max_seqlen,
+            max_seqlen,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+            block_table=None,
+        )
+        if is_grad:
+            ctx.save_for_backward(
+                q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state
+            )
+            ctx.dropout_p = dropout_p
+            ctx.max_seqlen = max_seqlen
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors
+        qkv_shape = q.shape[:-2] + (3, *q.shape[-2:])
+        dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device)
+        head_size_og = dout.size(2)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_varlen_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dqkv[:, 0],
+            dqkv[:, 1],
+            dqkv[:, 2],
+            cu_seqlens,
+            cu_seqlens,
+            ctx.max_seqlen,
+            ctx.max_seqlen,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dqkv = dqkv[..., : dout.shape[-1]]  # We could have padded the head dimension
+        return dqkv, None, None, None, None, None, None, None, None, None, None, None
+class FlashAttnKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        kv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach()
+        head_size_og = q.size(3)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+        )
+        if is_grad:
+            ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
+        dq = torch.empty_like(q)
+        kv_shape = k.shape[:-2] + (2, *k.shape[-2:])
+        dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device)
+        head_size_og = dout.size(3)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dkv[:, :, 0],
+            dkv[:, :, 1],
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dkv = dkv[..., : dout.shape[-1]]
+        return dq, dkv, None, None, None, None, None, None, None, None, None
+class FlashAttnVarlenKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        kv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        k, v = kv[:, 0].detach(), kv[:, 1].detach()
+        head_size_og = q.size(2)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+            block_table=None,
+        )
+        if is_grad:
+            ctx.save_for_backward(
+                q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state
+            )
+            ctx.dropout_p = dropout_p
+            ctx.max_seqlen_q = max_seqlen_q
+            ctx.max_seqlen_k = max_seqlen_k
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = (
+            ctx.saved_tensors
+        )
+        dq = torch.empty_like(q)
+        kv_shape = k.shape[:-2] + (2, *k.shape[-2:])
+        dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device)
+        head_size_og = dout.size(2)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_varlen_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dkv[:, 0],
+            dkv[:, 1],
+            cu_seqlens_q,
+            cu_seqlens_k,
+            ctx.max_seqlen_q,
+            ctx.max_seqlen_k,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dkv = dkv[..., : dout.shape[-1]]
+        return (
+            dq,
+            dkv,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+class FlashAttnFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        head_size_og = q.size(3)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+        )
+        if is_grad:
+            ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
+        dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+        head_size_og = dout.size(3)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dk,
+            dv,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dk = dk[..., : dout.shape[-1]]
+        dv = dv[..., : dout.shape[-1]]
+        return dq, dk, dv, None, None, None, None, None, None, None, None, None
+class FlashAttnVarlenFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        block_table,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        head_size_og = q.size(2)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+            block_table=block_table,
+        )
+        if is_grad:
+            ctx.save_for_backward(
+                q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state
+            )
+            ctx.dropout_p = dropout_p
+            ctx.max_seqlen_q = max_seqlen_q
+            ctx.max_seqlen_k = max_seqlen_k
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = (
+            ctx.saved_tensors
+        )
+        dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+        head_size_og = dout.size(2)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_varlen_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dk,
+            dv,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            ctx.max_seqlen_q,
+            ctx.max_seqlen_k,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dk = dk[..., : dout.shape[-1]]
+        dv = dv[..., : dout.shape[-1]]
+        return (
+            dq,
+            dk,
+            dv,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+def flash_attn_qkvpacked_func(
+    qkv,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # <=0.0 means deactivate
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If Q, K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of Q, K, V.
+    For multi-query and grouped-query attention (MQA/GQA), please see
+    flash_attn_kvpacked_func and flash_attn_func.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.
+    Arguments:
+        qkv: (batch_size, seqlen, 3, nheads, headdim)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to
+            the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnQKVPackedFunc.apply(
+        qkv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_kvpacked_func(
+    q,
+    kv,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of K, V.
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        kv: (batch_size, seqlen, 2, nheads_k, headdim)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnKVPackedFunc.apply(
+        q,
+        kv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_func(
+    q,
+    k,
+    v,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        k: (batch_size, seqlen, nheads_k, headdim)
+        v: (batch_size, seqlen, nheads_k, headdim)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnFunc.apply(
+        q,
+        k,
+        v,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_varlen_qkvpacked_func(
+    qkv,
+    cu_seqlens,
+    max_seqlen,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If Q, K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_varlen_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of Q, K, V.
+    For multi-query and grouped-query attention (MQA/GQA), please see
+    flash_attn_varlen_kvpacked_func and flash_attn_varlen_func.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.
+    Arguments:
+        qkv: (total, 3, nheads, headdim), where total = total number of tokens in the batch.
+        cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into qkv.
+        max_seqlen: int. Maximum sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnVarlenQKVPackedFunc.apply(
+        qkv,
+        cu_seqlens,
+        max_seqlen,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_varlen_kvpacked_func(
+    q,
+    kv,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of K, V.
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.
+        kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into q.
+        cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into kv.
+        max_seqlen_q: int. Maximum query sequence length in the batch.
+        max_seqlen_k: int. Maximum key sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnVarlenKVPackedFunc.apply(
+        q,
+        kv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_varlen_func(
+    q,
+    k,
+    v,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+    block_table=None,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.
+        k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into q.
+        cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into kv.
+        max_seqlen_q: int. Maximum query sequence length in the batch.
+        max_seqlen_k: int. Maximum key sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnVarlenFunc.apply(
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        block_table,
+        False if q.device.type == "cpu" else torch.is_grad_enabled(),
+    )
+def flash_attn_with_kvcache(
+    q,
+    k_cache,
+    v_cache,
+    k=None,
+    v=None,
+    rotary_cos=None,
+    rotary_sin=None,
+    cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None,
+    cache_batch_idx: Optional[torch.Tensor] = None,
+    cache_leftpad: Optional[torch.Tensor] = None,
+    block_table: Optional[torch.Tensor] = None,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    rotary_interleaved=True,
+    alibi_slopes=None,
+    num_splits=0,
+    return_softmax_lse=False,
+):
+    """
+    If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from
+    k and v. This is useful for incremental decoding: you can pass in the cached keys/values from
+    the previous step, and update them with the new keys/values from the current step, and do
+    attention with the updated cache, all in 1 kernel.
+    If you pass in k / v, you must make sure that the cache is large enough to hold the new values.
+    For example, the KV cache could be pre-allocated with the max sequence length, and you can use
+    cache_seqlens to keep track of the current sequence lengths of each sequence in the batch.
+    Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be
+    rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
+    If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos
+    and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
+    If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at
+    indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens).
+    See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function.
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Note: Does not support backward pass.
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,
+            or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)
+            page_block_size must be a multiple of 256.
+        v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,
+            or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)
+        k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate
+            k with k_cache, starting at the indices specified by cache_seqlens.
+        v [optional]: (batch_size, seqlen_new, nheads_k, headdim). Similar to k.
+        rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding
+            to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16.
+        rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos.
+        cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the
+            KV cache.
+        cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache.
+            If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1].
+            If the indices are not distinct, and k and v are provided, the values updated in the cache
+                 might come from any of the duplicate indices.
+        cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0.
+        block_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in.
+            If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False,
+            rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1
+            (i.e. GPT-NeoX style).
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        num_splits: int. If > 1, split the key/value into this many chunks along the sequence.
+           If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic
+           to automatically determine the number of splits.
+           Don't change this unless you know what you are doing.
+        return_softmax_lse: bool. Whether to return the logsumexp of the attention scores.
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+    """
+    assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension"
+    assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension"
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** (-0.5)
+    if cache_seqlens is not None and isinstance(cache_seqlens, int):
+        cache_seqlens = torch.full(
+            (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device
+        )
+        cache_seqlens = maybe_contiguous(cache_seqlens)
+    cache_batch_idx = maybe_contiguous(cache_batch_idx)
+    block_table = maybe_contiguous(block_table)
+    out, softmax_lse = flash_attn.fwd_kvcache(
+        q,
+        k_cache,
+        v_cache,
+        k,
+        v,
+        cache_seqlens,
+        rotary_cos,
+        rotary_sin,
+        cache_batch_idx,
+        cache_leftpad,
+        block_table,
+        alibi_slopes,
+        None,
+        softmax_scale,
+        causal,
+        window_size[0],
+        window_size[1],
+        softcap,
+        rotary_interleaved,
+        num_splits,
+    )
+    return (out, softmax_lse) if return_softmax_lse else out

build/torch-stable-abi210-cu132-aarch64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,39 @@

+{
+  "name": "flash-attn2",
+  "id": "_flash_attn2_cuda_460047c",
+  "version": 4,
+  "kernels-minver": "0.14.0",
+  "license": "BSD-3-Clause",
+  "python-depends": [],
+  "kernel-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Nz8hqrusjoUHQqszpxfF2USMwUXn1VfpvcmtAAvutSM=",
+      "_flash_attn2_cuda_460047c.abi3.so": "o4KfnnQPzHLqZ81TNKaReUO+g2+GvWUpV/1i8eOuo7k=",
+      "_ops.py": "DV4Eo4+s3aaNmDsfXoZiYKNLnRUJqlsdDGlMsdKFgk4=",
+      "flash_attn_interface.py": "6UYUhfUax4+yOUygyS7K8dVjrIjU5GQr1AJ2gSp8rPo=",
+      "symbols.json": "uvC+87tHUhtHv7L/Nl53T/Ms7ceX317/jwQBoqHY14A="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.18.0-dev0",
+      "commit": "d054906a13710b811b96697b4635f03020842733",
+      "dirty": false
+    },
+    "kernel": {
+      "commit": "460047c291fc7f0b3dbd4239b46bd5e1281b1806",
+      "dirty": false
+    }
+  }
+}

build/torch-stable-abi210-cu132-aarch64-linux/metadata.json.sigstore ADDED Viewed

	@@ -0,0 +1 @@

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHKzCCBrGgAwIBAgIUN7yA9ogS+KkeDYuoPcAUr6hsD4gwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDI0MDAwNzUyWhcNMjYxMDI0MDAxNzUyWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEMN5ZaLX/ccjUFKKfTQ5XP129/yHvuxFPcETq71aRGfd4lZOVwhPNa/yMC2rJxxdmvIlnXXnM71c9s8j+4rdsRqOCBdAwggXMMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBCg0NjAwNDdjMjkxZmM3ZjBiM2RiZDQyMzliNDZiZDVlMTI4MWIxODA2MBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDCg0NjAwNDdjMjkxZmM3ZjBiM2RiZDQyMzliNDZiZDVlMTI4MWIxODA2MBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDCg0NjAwNDdjMjkxZmM3ZjBiM2RiZDQyMzliNDZiZDVlMTI4MWIxODA2MB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoNDYwMDQ3YzI5MWZjN2YwYjNkYmQ0MjM5YjQ2YmQ1ZTEyODFiMTgwNjAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTk4MjY0MDUxNS9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGJBgorBgEEAdZ5AgQCBHsEeQB3AHUA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhazw+FwAABAMARjBEAiBhXXca/RPxAQD6x2U/fkkQs2Qb8xhB3Ocbtk8EAL85GwIgFLfhJk5UBvp5BCpQy306uX7zrE0ceQ6zVJ6sLk64CJ8wCgYIKoZIzj0EAwMDaAAwZQIwSu3h0shlf5zMVbi2esJOvdB3yo1tjfdSU3p0BnUQqTQGDcIbzFRw4daoStzhwG6sAjEAhZifr0fUQ3tKF7NYXYHSqfURaRZef+oxxMOp2on1bYztf6njLa8bQ5ofvAYRRyDu"},"tlogEntries":[{"logIndex":"3430232829","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792800472","inclusionPromise":{"signedEntryTimestamp":"MEUCIQDapAVpzW9ZMQuTUJ39O7fTsatpIEek0pompEwTH4wvswIgRVThrMi15G2rOZgieSQkQYoUfgxwA9jQoW3N99QqtOU="},"inclusionProof":{"logIndex":"3308328567","rootHash":"8d7Q7d3LmggWU8eU5xZvu5mR+vnArSvWFl8O3w51kQ4=","treeSize":"3308328619","hashes":["ZeSraJnEOIg5RprEttxTPyJoHEIpK2DU9f+GgF7gKUU=","cTgqACgBgsyVC/yK4IGTbMOaCKHRO6569hXM02AYU7Q=","mlwEM8vnsselj+orQh+ahYjAcQRH/p00zGXam6wJfpg=","Ky499GSNP9wO5tKXqV1xgq2eVFaO9tYJrSgCADxcRvY=","C/FFp0p++7+e4msx8G01Pjp5EsyyHe7ThgIj+TssXXQ=","/ubypAm95ilTKdYwhtn6ndwUwRpZ3ICLmEdRdWZB5c8=","gWAOr6FsG8QrENGVRNoJq5UmkHkoGgTYikl74JYz6mg=","eiXoi0AkFmtS6KaPVPsLZdJkcMd6L/c9gNFiiUO7E+g=","K/8DiKTfabbp+AGoyQMoqW/REHMadjMC9YuSpdz5/Cs=","8jJv+xlIJWDTCEJK91Tx1dYXJK/tgcxFsykheeCSw3U=","gdS39fsn5MOVlKluUgGY2sdKF3H5oauxkHMa7DSSNcg=","8x0N/oAmvpY+6DQUGTU+i6a1Cuqphro2WbV5WZ/RHL8=","DMchGrSAeKV9qHo75XJV3UFZ/CHF8zofPRu+NgVRwA4=","V6dXjeDcwYxBG/9gVSBM8rYdPrAqGpZsf1e7v6l67GA=","iDfImp65w8dirRfYGNgFA34JTUYDWbcJwNBBU/Ok+X4=","YYCr61A8azar3dfO8Ct+pri7LqDxkUlUkfs7fiqYErU=","1DaDKzrCNNBhmSdY5vFjz+NpAmxmYfYEgBCdPLrMP1A=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3308328619\n8d7Q7d3LmggWU8eU5xZvu5mR+vnArSvWFl8O3w51kQ4=\n\n— rekor.sigstore.dev wNI9ajBFAiEAl3psN46GEcjBpud4+y/WZ4F3HtMSO1IV8ZardGaTWcACIEjD4ZIUStqCf5b+jEh+/EbgTp70rKCnzpq9mNBPkHLQ\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIyNWMyZDk2MDQwNzkxYTJhNmMxZjE4ZDg0MjkwNDQxOTJlY2JjZTMzZTNiMjg4YjM1MDUyYWYxNDRjNGVlM2QzIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJRkJXa2hMNDl4V09YY3Vwdmh3VlBFZk14RWxPR0tvVmp1azlxZjF6OW1GMUFpRUF5Q0VaLzVTNmhENGNEajcwOUszUXl3elRmSUo1UkF6S1MrbjFNZ3RLTGtvPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaExla05EUW5KSFowRjNTVUpCWjBsVlRqZDVRVGx2WjFNclMydGxSRmwxYjFCalFWVnlObWh6UkRSbmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFU1RCTlJFRjNUbnBWZVZkb1kwNU5hbGw0VFVSSk1FMUVRWGhPZWxWNVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZOVGpWYVlVeFlMMk5qYWxWR1MwdG1WRkUxV0ZBeE1qa3ZlVWgyZFhoR1VHTkZWSEVLTnpGaFVrZG1aRFJzV2s5V2QyaFFUbUV2ZVUxRE1uSktlSGhrYlhaSmJHNVlXRzVOTnpGak9YTTRhaXMwY21SelVuRlBRMEprUVhkbloxaE5UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmN3VG1wQmR3cE9SR1JxVFdwcmVGcHRUVE5hYWtKcFRUSlNhVnBFVVhsTmVteHBUa1JhYVZwRVZteE5WRWswVFZkSmVFOUVRVEpOUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEWnpCT2FrRjNUa1JrYWsxcWEzaGFiVTB6Q2xwcVFtbE5NbEpwV2tSUmVVMTZiR2xPUkZwcFdrUldiRTFVU1RSTlYwbDRUMFJCTWsxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEWnpCT2FrRjNDazVFWkdwTmFtdDRXbTFOTTFwcVFtbE5NbEpwV2tSUmVVMTZiR2xPUkZwcFdrUldiRTFVU1RSTlYwbDRUMFJCTWsxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlPUkZsM1RVUlJNMWw2U1RWTlYxcHFUakpaZDFscVRtdFpiVkV3VFdwTk5WbHFVVEpaYlZFeFdsUkZlVTlFUm1rS1RWUm5kMDVxUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUnJORTFxV1RCTlJGVjRUbE01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkS0NrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NITkZaVkZDTTBGSVZVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvWVhwM0swWjNRVUZDUVUxQlVtcENSVUZwUW1oWVdHTmhMMUpRZUVGUlJEWjRNbFV2Wm10clVYTXlVV0k0ZUdoQ00wOWpZZ3AwYXpoRlFVdzROVWQzU1dkR1RHWm9TbXMxVlVKMmNEVkNRM0JSZVRNd05uVllOM3B5UlRCalpWRTJlbFpLTm5OTWF6WTBRMG80ZDBObldVbExiMXBKQ25wcU1FVkJkMDFFWVVGQmQxcFJTWGRUZFROb01ITm9iR1kxZWsxV1lta3laWE5LVDNaa1FqTjViekYwYW1aa1UxVXpjREJDYmxWUmNWUlJSMFJqU1dJS2VrWlNkelJrWVc5VGRIcG9kMGMyYzBGcVJVRm9XbWxtY2pCbVZWRXpkRXRHTjA1WldGbElVM0ZtVlZKaFVscGxaaXR2ZUhoTlQzQXliMjR4WWxsNmRBcG1ObTVxVEdFNFlsRTFiMloyUVZsU1VubEVkUW90TFMwdExVVk9SQ0JEUlZKVVNVWkpRMEZVUlMwdExTMHRDZz09In19fX0="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyDADAgEAMIICvwYJKoZIhvcNAQcCoIICsDCCAqwCAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgPBaos2ErBxosU60dEUHP72xA3r2tJxcZ+U5pO+kAzzgCFFY/daZCnGNuqdgKsimEVKQXAxKiGA8yMDI2MTAyNDAwMDc1MlowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdowggHWAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDI0MDAwNzUyWjAvBgkqhkiG9w0BCQQxIgQgDWv5Ubsq3b4BUyPtOKBjEHmam6qOnXf0Ducy0Q/ALagwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGYwZAIxAMGn96mgL9G4S6v8KAhsXgaLHLxLtLXkMEuVxlrkSey5PYJmitundFbhChOAUAn5UQIvVczNzLLzW54p+k/liiht08kDHlfLylc4CehlBIyOx4XANNwmxu4LW0fhJxYKjNI="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"JcLZYEB5GipsHxjYQpBEGS7LzjPjsoizUFKvFExO49M="},"signature":"MEUCIFBWkhL49xWOXcupvhwVPEfMxElOGKoVjuk9qf1z9mF1AiEAyCEZ/5S6hD4cDj709K3QywzTfIJ5RAzKS+n1MgtKLko="}}

build/torch-stable-abi210-cu132-aarch64-linux/symbols.json ADDED Viewed

	@@ -0,0 +1,1273 @@

+{
+  "schema_version": 1,
+  "module": "_flash_attn2_cuda_460047c",
+  "functions": [
+    {
+      "name": "flash_attn_func",
+      "qualname": "flash_attn_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (batch_size, seqlen, nheads, headdim)\n    k: (batch_size, seqlen, nheads_k, headdim)\n    v: (batch_size, seqlen, nheads_k, headdim)\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_kvpacked_func",
+      "qualname": "flash_attn_kvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of K, V.\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (batch_size, seqlen, nheads, headdim)\n    kv: (batch_size, seqlen, 2, nheads_k, headdim)\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "kv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_qkvpacked_func",
+      "qualname": "flash_attn_qkvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf Q, K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of Q, K, V.\nFor multi-query and grouped-query attention (MQA/GQA), please see\nflash_attn_kvpacked_func and flash_attn_func.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.\n\nArguments:\n    qkv: (batch_size, seqlen, 3, nheads, headdim)\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to\n        the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "qkv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_varlen_func",
+      "qualname": "flash_attn_varlen_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.\n    k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.\n    v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.\n    cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into q.\n    cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into kv.\n    max_seqlen_q: int. Maximum query sequence length in the batch.\n    max_seqlen_k: int. Maximum key sequence length in the batch.\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (total, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_varlen_kvpacked_func",
+      "qualname": "flash_attn_varlen_kvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of K, V.\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.\n    kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch.\n    cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into q.\n    cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into kv.\n    max_seqlen_q: int. Maximum query sequence length in the batch.\n    max_seqlen_k: int. Maximum key sequence length in the batch.\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (total, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "kv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_varlen_qkvpacked_func",
+      "qualname": "flash_attn_varlen_qkvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf Q, K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_varlen_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of Q, K, V.\nFor multi-query and grouped-query attention (MQA/GQA), please see\nflash_attn_varlen_kvpacked_func and flash_attn_varlen_func.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.\n\nArguments:\n    qkv: (total, 3, nheads, headdim), where total = total number of tokens in the batch.\n    cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into qkv.\n    max_seqlen: int. Maximum sequence length in the batch.\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (total, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "qkv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_with_kvcache",
+      "qualname": "flash_attn_with_kvcache",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from\nk and v. This is useful for incremental decoding: you can pass in the cached keys/values from\nthe previous step, and update them with the new keys/values from the current step, and do\nattention with the updated cache, all in 1 kernel.\n\nIf you pass in k / v, you must make sure that the cache is large enough to hold the new values.\nFor example, the KV cache could be pre-allocated with the max sequence length, and you can use\ncache_seqlens to keep track of the current sequence lengths of each sequence in the batch.\n\nAlso apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be\nrotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.\nIf causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos\nand rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.\nIf not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at\nindices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens).\n\nSee tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function.\n\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nNote: Does not support backward pass.\n\nArguments:\n    q: (batch_size, seqlen, nheads, headdim)\n    k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,\n        or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)\n        page_block_size must be a multiple of 256.\n    v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,\n        or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)\n    k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate\n        k with k_cache, starting at the indices specified by cache_seqlens.\n    v [optional]: (batch_size, seqlen_new, nheads_k, headdim). Similar to k.\n    rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding\n        to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16.\n    rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos.\n    cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the\n        KV cache.\n    cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache.\n        If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1].\n        If the indices are not distinct, and k and v are provided, the values updated in the cache\n             might come from any of the duplicate indices.\n    cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0.\n    block_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in.\n        If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False,\n        rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1\n        (i.e. GPT-NeoX style).\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    num_splits: int. If > 1, split the key/value into this many chunks along the sequence.\n       If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic\n       to automatically determine the number of splits.\n       Don't change this unless you know what you are doing.\n    return_softmax_lse: bool. Whether to return the logsumexp of the attention scores.\n\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k_cache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "v_cache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "rotary_cos",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "rotary_sin",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "cache_seqlens",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "int | torch.Tensor | None"
+          },
+          {
+            "name": "cache_batch_idx",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "cache_leftpad",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "rotary_interleaved",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "True",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "num_splits",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": null
+          },
+          {
+            "name": "return_softmax_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "fwd",
+      "qualname": "fwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Forward pass for multi-head attention.\n\nArgs:\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    out: Optional output tensor, same shape as q\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    return_softmax: Whether to return softmax weights\n    gen: Optional random number generator\n\nReturns:\n    List of tensors: [output, softmax_lse, (softmax if return_softmax)]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "return_softmax",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "varlen_fwd",
+      "qualname": "varlen_fwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Forward pass for multi-head attention with variable sequence lengths.\n\nArgs:\n    q: Query tensor of shape [total_q, num_heads, head_size]\n    k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]\n    cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]\n    out: Optional output tensor of shape [total_q, num_heads, head_size]\n    seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size]\n    leftpad_k: Optional left padding for keys of shape [batch_size]\n    block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    max_seqlen_q: Maximum sequence length for queries\n    max_seqlen_k: Maximum sequence length for keys\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    zero_tensors: Whether to zero tensors before computation\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    return_softmax: Whether to return softmax weights\n    gen: Optional random number generator\n\nReturns:\n    List of tensors: [output, softmax_lse, (softmax if return_softmax)]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "seqused_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "leftpad_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "zero_tensors",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "return_softmax",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "bwd",
+      "qualname": "bwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Backward pass for multi-head attention.\n\nArgs:\n    dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]\n    softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]\n    dq: Optional gradient tensor for queries, same shape as q\n    dk: Optional gradient tensor for keys, same shape as k\n    dv: Optional gradient tensor for values, same shape as v\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    deterministic: Whether to use deterministic algorithms\n    gen: Optional random number generator\n    rng_state: Optional RNG state from forward pass\n\nReturns:\n    List of tensors: [dq, dk, dv]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "dout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "softmax_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "dq",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dk",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          },
+          {
+            "name": "rng_state",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "varlen_bwd",
+      "qualname": "varlen_bwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Backward pass for multi-head attention with variable sequence lengths.\n\nArgs:\n    dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]\n    softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]\n    cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]\n    cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]\n    dq: Optional gradient tensor for queries, same shape as q\n    dk: Optional gradient tensor for keys, same shape as k\n    dv: Optional gradient tensor for values, same shape as v\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    max_seqlen_q: Maximum sequence length for queries\n    max_seqlen_k: Maximum sequence length for keys\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    zero_tensors: Whether to zero tensors before computation\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    deterministic: Whether to use deterministic algorithms\n    gen: Optional random number generator\n    rng_state: Optional RNG state from forward pass\n\nReturns:\n    List of tensors: [dq, dk, dv]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "dout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "softmax_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "dq",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dk",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "zero_tensors",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          },
+          {
+            "name": "rng_state",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "fwd_kvcache",
+      "qualname": "fwd_kvcache",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Forward pass for multi-head attention with KV cache.\n\nArgs:\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]\n    v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]\n    seqlens_k: Optional sequence lengths for keys of shape [batch_size]\n    rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2]\n    rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2]\n    cache_batch_idx: Optional indices to index into the KV cache\n    leftpad_k: Optional left padding for keys of shape [batch_size]\n    block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    out: Optional output tensor, same shape as q\n    softmax_scale: Scale factor for softmax\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    is_rotary_interleaved: Whether rotary embeddings are interleaved\n    num_splits: Number of splits for computation\n\nReturns:\n    List of tensors: [output, softmax_lse]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "kcache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "vcache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "rotary_cos",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "rotary_sin",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "cache_batch_idx",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "leftpad_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "is_rotary_interleaved",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "num_splits",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "1",
+            "annotation": "int"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    }
+  ],
+  "layers": []
+}

build/torch-stable-abi210-cu134-aarch64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,424 @@

+from typing import List, Optional
+import torch
+from . import flash_attn_interface
+from ._ops import ops as flash_attn_ops
+from .flash_attn_interface import (
+    # helper to set up the Philox random number generator for dropout
+    _philox_generator,
+    # reexported for convenience to be used in the high-level API functions
+    flash_attn_func,
+    flash_attn_kvpacked_func,
+    flash_attn_qkvpacked_func,
+    flash_attn_varlen_func,
+    flash_attn_varlen_kvpacked_func,
+    flash_attn_varlen_qkvpacked_func,
+    flash_attn_with_kvcache,
+)
+__all__ = [
+    # High-level API (with autograd support)
+    "flash_attn_func",
+    "flash_attn_kvpacked_func",
+    "flash_attn_qkvpacked_func",
+    "flash_attn_varlen_func",
+    "flash_attn_varlen_kvpacked_func",
+    "flash_attn_varlen_qkvpacked_func",
+    "flash_attn_with_kvcache",
+    # Upstream interface module. Its `_flash_attn_forward`, `_flash_attn_backward`,
+    # `_flash_attn_varlen_forward` and `_flash_attn_varlen_backward` are supported
+    # low-level entry points with upstream's signatures: Torch custom ops without
+    # autograd that return the LSE and are traceable by torch.compile. The autograd
+    # Functions above are built on top of them.
+    "flash_attn_interface",
+    # Low-level ops (no autograd)
+    "fwd",
+    "varlen_fwd",
+    "bwd",
+    "varlen_bwd",
+    "fwd_kvcache",
+]
+def fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    return_softmax: bool = False,
+    gen: Optional[torch.Generator] = None,
+) -> List[torch.Tensor]:
+    """
+    Forward pass for multi-head attention.
+    Args:
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        out: Optional output tensor, same shape as q
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        return_softmax: Whether to return softmax weights
+        gen: Optional random number generator
+    Returns:
+        List of tensors: [output, softmax_lse, (softmax if return_softmax)]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    gen = _philox_generator(q, p_dropout, gen)
+    return flash_attn_ops.fwd(
+        q,
+        k,
+        v,
+        out,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+def varlen_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    block_table: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    max_seqlen_q: int = 0,
+    max_seqlen_k: int = 0,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    zero_tensors: bool = False,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    return_softmax: bool = False,
+    gen: Optional[torch.Generator] = None,
+) -> List[torch.Tensor]:
+    """
+    Forward pass for multi-head attention with variable sequence lengths.
+    Args:
+        q: Query tensor of shape [total_q, num_heads, head_size]
+        k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]
+        cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]
+        out: Optional output tensor of shape [total_q, num_heads, head_size]
+        seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size]
+        leftpad_k: Optional left padding for keys of shape [batch_size]
+        block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        max_seqlen_q: Maximum sequence length for queries
+        max_seqlen_k: Maximum sequence length for keys
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        zero_tensors: Whether to zero tensors before computation
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        return_softmax: Whether to return softmax weights
+        gen: Optional random number generator
+    Returns:
+        List of tensors: [output, softmax_lse, (softmax if return_softmax)]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    gen = _philox_generator(q, p_dropout, gen)
+    return flash_attn_ops.varlen_fwd(
+        q,
+        k,
+        v,
+        out,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_k,
+        leftpad_k,
+        block_table,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        p_dropout,
+        softmax_scale,
+        zero_tensors,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+def bwd(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor] = None,
+    dk: Optional[torch.Tensor] = None,
+    dv: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    deterministic: bool = False,
+    gen: Optional[torch.Generator] = None,
+    rng_state: Optional[torch.Tensor] = None,
+) -> List[torch.Tensor]:
+    """
+    Backward pass for multi-head attention.
+    Args:
+        dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]
+        softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]
+        dq: Optional gradient tensor for queries, same shape as q
+        dk: Optional gradient tensor for keys, same shape as k
+        dv: Optional gradient tensor for values, same shape as v
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        deterministic: Whether to use deterministic algorithms
+        gen: Optional random number generator
+        rng_state: Optional RNG state from forward pass
+    Returns:
+        List of tensors: [dq, dk, dv]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    return flash_attn_ops.bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        alibi_slopes,
+        p_dropout,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        gen,
+        rng_state,
+    )
+def varlen_bwd(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    dq: Optional[torch.Tensor] = None,
+    dk: Optional[torch.Tensor] = None,
+    dv: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    max_seqlen_q: int = 0,
+    max_seqlen_k: int = 0,
+    p_dropout: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    zero_tensors: bool = False,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    deterministic: bool = False,
+    gen: Optional[torch.Generator] = None,
+    rng_state: Optional[torch.Tensor] = None,
+) -> List[torch.Tensor]:
+    """
+    Backward pass for multi-head attention with variable sequence lengths.
+    Args:
+        dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]
+        out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]
+        softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]
+        cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]
+        cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]
+        dq: Optional gradient tensor for queries, same shape as q
+        dk: Optional gradient tensor for keys, same shape as k
+        dv: Optional gradient tensor for values, same shape as v
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        max_seqlen_q: Maximum sequence length for queries
+        max_seqlen_k: Maximum sequence length for keys
+        p_dropout: Dropout probability
+        softmax_scale: Scale factor for softmax
+        zero_tensors: Whether to zero tensors before computation
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        deterministic: Whether to use deterministic algorithms
+        gen: Optional random number generator
+        rng_state: Optional RNG state from forward pass
+    Returns:
+        List of tensors: [dq, dk, dv]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    return flash_attn_ops.varlen_bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        p_dropout,
+        softmax_scale,
+        zero_tensors,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        gen,
+        rng_state,
+    )
+def fwd_kvcache(
+    q: torch.Tensor,
+    kcache: torch.Tensor,
+    vcache: torch.Tensor,
+    k: Optional[torch.Tensor] = None,
+    v: Optional[torch.Tensor] = None,
+    seqlens_k: Optional[torch.Tensor] = None,
+    rotary_cos: Optional[torch.Tensor] = None,
+    rotary_sin: Optional[torch.Tensor] = None,
+    cache_batch_idx: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    block_table: Optional[torch.Tensor] = None,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    is_causal: bool = False,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    is_rotary_interleaved: bool = False,
+    num_splits: int = 1,
+) -> List[torch.Tensor]:
+    """
+    Forward pass for multi-head attention with KV cache.
+    Args:
+        q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]
+        kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]
+        k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]
+        v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]
+        seqlens_k: Optional sequence lengths for keys of shape [batch_size]
+        rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2]
+        rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2]
+        cache_batch_idx: Optional indices to index into the KV cache
+        leftpad_k: Optional left padding for keys of shape [batch_size]
+        block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]
+        alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]
+        out: Optional output tensor, same shape as q
+        softmax_scale: Scale factor for softmax
+        is_causal: Whether to use causal attention
+        window_size_left: Window size for left context (-1 for unlimited)
+        window_size_right: Window size for right context (-1 for unlimited)
+        softcap: Soft cap for attention weights
+        is_rotary_interleaved: Whether rotary embeddings are interleaved
+        num_splits: Number of splits for computation
+    Returns:
+        List of tensors: [output, softmax_lse]
+    """
+    if softmax_scale is None:
+        attention_head_dim = q.shape[-1]
+        softmax_scale = 1.0 / (attention_head_dim**0.5)
+    return flash_attn_ops.fwd_kvcache(
+        q,
+        kcache,
+        vcache,
+        k,
+        v,
+        seqlens_k,
+        rotary_cos,
+        rotary_sin,
+        cache_batch_idx,
+        leftpad_k,
+        block_table,
+        alibi_slopes,
+        out,
+        softmax_scale,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        is_rotary_interleaved,
+        num_splits,
+    )

build/torch-stable-abi210-cu134-aarch64-linux/_flash_attn2_cuda_460047c.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:923bd50971a3d94f45b9a5fabc2223a7fdca1d40d8ccff4848907eb850e99b4c
+size 1011746240

build/torch-stable-abi210-cu134-aarch64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _flash_attn2_cuda_460047c
+ops = torch.ops._flash_attn2_cuda_460047c
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_flash_attn2_cuda_460047c::{op_name}"

build/torch-stable-abi210-cu134-aarch64-linux/flash_attn_interface.py ADDED Viewed

	@@ -0,0 +1,1708 @@

+# Copyright (c) 2023, Tri Dao.
+from typing import Optional, Sequence, Tuple, Union
+import torch
+import torch.nn as nn
+import os
+# # isort: off
+# # We need to import the CUDA kernels after importing torch
+# USE_TRITON_ROCM = os.getenv("FLASH_ATTENTION_TRITON_AMD_ENABLE", "FALSE") == "TRUE"
+# if USE_TRITON_ROCM:
+#     from .flash_attn_triton_amd import interface_fa as flash_attn
+# else:
+#     import flash_attn_2_cuda as flash_attn
+from ._ops import ops as flash_attn
+from ._ops import add_op_namespace_prefix
+# # isort: on
+# Torch >= 2.13 can pass a Generator through the stable ABI. Older versions reject the argument
+# before the kernel runs, so there the CUDA op has to draw its Philox seed and offset itself.
+_STABLE_ABI_GENERATOR = torch.__version__ >= "2.13"
+def maybe_contiguous(x):
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+def _get_device():
+    if torch.xpu.is_available():
+        return "xpu"
+    elif torch.cuda.is_available():
+        return "cuda"
+    else:
+        return "cpu"
+def _get_block_size_n(device, head_dim, is_dropout, is_causal):
+    assert head_dim <= 256
+    if device.type == "xpu":
+        if head_dim <= 96:
+            return 64
+        elif head_dim <= 128:
+            return 32
+        elif head_dim <= 256:
+            return 64
+        else:
+            return 32
+    # This should match the block sizes in the CUDA kernel
+    major, minor = torch.cuda.get_device_capability(device)
+    is_sm8x = (
+        major == 8 and minor > 0
+    )  # Only include sm86 and sm89, exclude sm80 (A100)
+    is_sm80 = major == 8 and minor == 0
+    is_sm90 = major == 9 and minor == 0
+    if head_dim <= 32:
+        return 128
+    if head_dim <= 64:
+        return 128 if not is_dropout else 64
+    elif head_dim <= 96:
+        return 64
+    elif head_dim <= 128:
+        if is_sm8x:
+            return 64 if (not is_dropout and is_causal) else 32
+        else:
+            return 64 if not is_dropout else 32
+    elif head_dim <= 192:
+        return 64
+    elif head_dim <= 224:
+        return 64
+    elif head_dim <= 256:
+        return 64
+def round_multiple(x, m):
+    return (x + m - 1) // m * m
+def _philox_generator(
+    q: torch.Tensor, dropout_p: float, generator: Optional[torch.Generator] = None
+) -> Optional[torch.Generator]:
+    """The generator to hand to `fwd` / `varlen_fwd`.
+    With dropout on CUDA this is the default generator unless one is given, so the kernel reserves
+    its seed and offset exactly like the non-stable build (see set_params_philox in flash_api.cpp).
+    """
+    if not q.is_cuda:
+        return generator
+    if dropout_p <= 0:
+        return None
+    if not _STABLE_ABI_GENERATOR:
+        if generator is not None:
+            raise ValueError(
+                "flash_attn2: passing a generator to the CUDA kernel requires torch >= 2.13; "
+                "seed the default generator with torch.manual_seed instead"
+            )
+        return None
+    if generator is None:
+        generator = torch.cuda.default_generators[q.device.index]
+    return generator
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_forward"),
+    mutates_args=(),
+    device_types=_get_device(),
+    # Tells torch.compile the forward draws random numbers, so it is never deduplicated or
+    # recomputed (e.g. under activation checkpointing).
+    tags=(torch.Tag.nondeterministic_seeded,),
+)
+def _flash_attn_forward(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    return_softmax: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    gen = _philox_generator(q, dropout_p)
+    out, softmax_lse, S_dmask, rng_state = flash_attn.fwd(
+        q,
+        k,
+        v,
+        None,
+        alibi_slopes,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    return out, softmax_lse, S_dmask, rng_state
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_forward"))
+def _flash_attn_forward_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    return_softmax: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    batch_size, seqlen_q, num_heads, head_size = q.shape
+    seqlen_k = k.shape[1]
+    out = torch.empty_like(q)
+    softmax_lse = torch.empty(
+        (batch_size, num_heads, seqlen_q),
+        dtype=torch.float32,
+        device=q.device,
+        layout=q.layout,
+    )
+    p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout)
+    if return_softmax:
+        p = torch.empty(
+            (
+                batch_size,
+                num_heads,
+                round_multiple(seqlen_q, 128),
+                round_multiple(seqlen_k, 128),
+            ),
+            dtype=q.dtype,
+            device=q.device,
+            layout=q.layout,
+        )
+    rng_state = torch.empty((2,), dtype=torch.int64, device=q.device)
+    return out, softmax_lse, p, rng_state
+_wrapped_flash_attn_forward = flash_attn._flash_attn_forward  # upstream alias, see CARD.md
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_varlen_forward"),
+    mutates_args=(),
+    device_types=_get_device(),
+    tags=(torch.Tag.nondeterministic_seeded,),
+)
+def _flash_attn_varlen_forward(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    return_softmax: bool = False,
+    block_table: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    gen = _philox_generator(q, dropout_p)
+    out, softmax_lse, S_dmask, rng_state = flash_attn.varlen_fwd(
+        q,
+        k,
+        v,
+        None,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_k,
+        leftpad_k,
+        block_table,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        zero_tensors,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        return_softmax,
+        gen,
+    )
+    # if out.isnan().any() or softmax_lse.isnan().any():
+    #     breakpoint()
+    return out, softmax_lse, S_dmask, rng_state
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_forward"))
+def _flash_attn_varlen_forward_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int = -1,
+    window_size_right: int = -1,
+    softcap: float = 0.0,
+    alibi_slopes: Optional[torch.Tensor] = None,
+    return_softmax: bool = False,
+    block_table: Optional[torch.Tensor] = None,
+    leftpad_k: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    paged_kv = block_table is not None
+    batch_size = cu_seqlens_q.numel() - 1
+    total_q, num_heads, _ = q.shape
+    out = torch.empty_like(q)
+    softmax_lse = torch.empty(
+        (num_heads, total_q), dtype=torch.float32, device=q.device, layout=q.layout
+    )
+    p = torch.empty((0,), dtype=q.dtype, device=q.device, layout=q.layout)
+    seqlen_q_rounded = round_multiple(max_seqlen_q, 128)
+    seqlen_k_rounded = round_multiple(max_seqlen_k, 128)
+    if return_softmax:
+        p = torch.empty(
+            (batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded),
+            dtype=q.dtype,
+            device=q.device,
+            layout=q.layout,
+        )
+    rng_state = torch.empty((2,), dtype=torch.int64, device=q.device)
+    return out, softmax_lse, p, rng_state
+_wrapped_flash_attn_varlen_forward = flash_attn._flash_attn_varlen_forward  # upstream alias, see CARD.md
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_backward"),
+    mutates_args=("dq", "dk", "dv"),
+    device_types=_get_device(),
+)
+def _flash_attn_backward(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    # dq, dk, dv are allocated by us so they should already be contiguous
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    (
+        dq,
+        dk,
+        dv,
+        softmax_d,
+    ) = flash_attn.bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        alibi_slopes,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        None,
+        rng_state,
+    )
+    return softmax_d
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_backward"))
+def _flash_attn_backward_fake(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    if dq is None:
+        dq = torch.empty_like(q)
+    if dk is None:
+        dk = torch.empty_like(k)
+    if dv is None:
+        dv = torch.empty_like(v)
+    batch_size, seqlen_q, num_heads, _ = q.shape
+    softmax_d = torch.empty(
+        (batch_size, num_heads, round_multiple(seqlen_q, 128)),
+        device=q.device,
+        dtype=torch.float32,
+    )
+    return softmax_d
+_wrapped_flash_attn_backward = flash_attn._flash_attn_backward  # upstream alias, see CARD.md
+@torch.library.custom_op(
+    add_op_namespace_prefix("_flash_attn_varlen_backward"),
+    mutates_args=("dq", "dk", "dv"),
+    device_types=_get_device(),
+)
+def _flash_attn_varlen_backward(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> torch.Tensor:
+    # dq, dk, dv are allocated by us so they should already be contiguous
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    (
+        dq,
+        dk,
+        dv,
+        softmax_d,
+    ) = flash_attn.varlen_bwd(
+        dout,
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        dq,
+        dk,
+        dv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        zero_tensors,
+        causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        deterministic,
+        None,
+        rng_state,
+    )
+    # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any():
+    #     breakpoint()
+    return softmax_d
+@torch.library.register_fake(add_op_namespace_prefix("_flash_attn_varlen_backward"))
+def _flash_attn_varlen_backward_fake(
+    dout: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    dq: Optional[torch.Tensor],
+    dk: Optional[torch.Tensor],
+    dv: Optional[torch.Tensor],
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    dropout_p: float,
+    softmax_scale: float,
+    causal: bool,
+    window_size_left: int,
+    window_size_right: int,
+    softcap: float,
+    alibi_slopes: Optional[torch.Tensor],
+    deterministic: bool,
+    rng_state: Optional[torch.Tensor] = None,
+    zero_tensors: bool = False,
+) -> torch.Tensor:
+    dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
+    batch_size = cu_seqlens_q.numel() - 1
+    total_q, num_heads, _ = q.shape
+    if dq is None:
+        dq = torch.empty_like(q)
+    if dk is None:
+        dk = torch.empty_like(k)
+    if dv is None:
+        dv = torch.empty_like(v)
+    softmax_d = torch.empty(
+        (num_heads, total_q + 128 * batch_size), device=q.device, dtype=torch.float32
+    )
+    return softmax_d
+_wrapped_flash_attn_varlen_backward = flash_attn._flash_attn_varlen_backward  # upstream alias, see CARD.md
+class FlashAttnQKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        qkv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and qkv.requires_grad
+        if softmax_scale is None:
+            softmax_scale = qkv.shape[-1] ** (-0.5)
+        q, k, v = qkv[:, :, 0].detach(), qkv[:, :, 1].detach(), qkv[:, :, 2].detach()
+        head_size_og = q.size(3)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+        )
+        if is_grad:
+            ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
+        qkv_shape = q.shape[:-2] + (3, *q.shape[-2:])
+        dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device)
+        head_size_og = dout.size(3)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dqkv[:, :, 0],
+            dqkv[:, :, 1],
+            dqkv[:, :, 2],
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dqkv = dqkv[..., : dout.shape[-1]]  # We could have padded the head dimension
+        return dqkv, None, None, None, None, None, None, None, None, None
+class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        qkv,
+        cu_seqlens,
+        max_seqlen,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and qkv.requires_grad
+        if softmax_scale is None:
+            softmax_scale = qkv.shape[-1] ** (-0.5)
+        q, k, v = qkv[:, 0].detach(), qkv[:, 1].detach(), qkv[:, 2].detach()
+        head_size_og = q.size(2)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
+            q,
+            k,
+            v,
+            cu_seqlens,
+            cu_seqlens,
+            max_seqlen,
+            max_seqlen,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+            block_table=None,
+        )
+        if is_grad:
+            ctx.save_for_backward(
+                q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state
+            )
+            ctx.dropout_p = dropout_p
+            ctx.max_seqlen = max_seqlen
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors
+        qkv_shape = q.shape[:-2] + (3, *q.shape[-2:])
+        dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device)
+        head_size_og = dout.size(2)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_varlen_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dqkv[:, 0],
+            dqkv[:, 1],
+            dqkv[:, 2],
+            cu_seqlens,
+            cu_seqlens,
+            ctx.max_seqlen,
+            ctx.max_seqlen,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dqkv = dqkv[..., : dout.shape[-1]]  # We could have padded the head dimension
+        return dqkv, None, None, None, None, None, None, None, None, None, None, None
+class FlashAttnKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        kv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        k, v = kv[:, :, 0].detach(), kv[:, :, 1].detach()
+        head_size_og = q.size(3)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+        )
+        if is_grad:
+            ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
+        dq = torch.empty_like(q)
+        kv_shape = k.shape[:-2] + (2, *k.shape[-2:])
+        dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device)
+        head_size_og = dout.size(3)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dkv[:, :, 0],
+            dkv[:, :, 1],
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dkv = dkv[..., : dout.shape[-1]]
+        return dq, dkv, None, None, None, None, None, None, None, None, None
+class FlashAttnVarlenKVPackedFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        kv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, kv])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        k, v = kv[:, 0].detach(), kv[:, 1].detach()
+        head_size_og = q.size(2)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+            block_table=None,
+        )
+        if is_grad:
+            ctx.save_for_backward(
+                q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state
+            )
+            ctx.dropout_p = dropout_p
+            ctx.max_seqlen_q = max_seqlen_q
+            ctx.max_seqlen_k = max_seqlen_k
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = (
+            ctx.saved_tensors
+        )
+        dq = torch.empty_like(q)
+        kv_shape = k.shape[:-2] + (2, *k.shape[-2:])
+        dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device)
+        head_size_og = dout.size(2)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_varlen_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dkv[:, 0],
+            dkv[:, 1],
+            cu_seqlens_q,
+            cu_seqlens_k,
+            ctx.max_seqlen_q,
+            ctx.max_seqlen_k,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dkv = dkv[..., : dout.shape[-1]]
+        return (
+            dq,
+            dkv,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+class FlashAttnFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        head_size_og = q.size(3)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_forward(
+            q,
+            k,
+            v,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+        )
+        if is_grad:
+            ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
+            ctx.dropout_p = dropout_p
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
+        dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+        head_size_og = dout.size(3)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dk,
+            dv,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dk = dk[..., : dout.shape[-1]]
+        dv = dv[..., : dout.shape[-1]]
+        return dq, dk, dv, None, None, None, None, None, None, None, None, None
+class FlashAttnVarlenFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_softmax,
+        block_table,
+        is_grad_enabled,
+    ):
+        is_grad = is_grad_enabled and any(x.requires_grad for x in [q, k, v])
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** (-0.5)
+        head_size_og = q.size(2)
+        if head_size_og % 8 != 0:
+            q = torch.nn.functional.pad(q, [0, 8 - head_size_og % 8])
+            k = torch.nn.functional.pad(k, [0, 8 - head_size_og % 8])
+            v = torch.nn.functional.pad(v, [0, 8 - head_size_og % 8])
+        out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            dropout_p,
+            softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            softcap=softcap,
+            alibi_slopes=alibi_slopes,
+            return_softmax=return_softmax and dropout_p > 0,
+            block_table=block_table,
+        )
+        if is_grad:
+            ctx.save_for_backward(
+                q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state
+            )
+            ctx.dropout_p = dropout_p
+            ctx.max_seqlen_q = max_seqlen_q
+            ctx.max_seqlen_k = max_seqlen_k
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.alibi_slopes = alibi_slopes
+            ctx.deterministic = deterministic
+        out = out_padded[..., :head_size_og]
+        return out if not return_softmax else (out, softmax_lse, S_dmask)
+    @staticmethod
+    def backward(ctx, dout, *args):
+        q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = (
+            ctx.saved_tensors
+        )
+        dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+        head_size_og = dout.size(2)
+        dout_padded = dout
+        if head_size_og % 8 != 0:
+            dout_padded = torch.nn.functional.pad(dout, [0, 8 - head_size_og % 8])
+        _flash_attn_varlen_backward(
+            dout_padded,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dk,
+            dv,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            ctx.max_seqlen_q,
+            ctx.max_seqlen_k,
+            ctx.dropout_p,
+            ctx.softmax_scale,
+            ctx.causal,
+            ctx.window_size[0],
+            ctx.window_size[1],
+            ctx.softcap,
+            ctx.alibi_slopes,
+            ctx.deterministic,
+            rng_state=rng_state,
+        )
+        dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
+        dk = dk[..., : dout.shape[-1]]
+        dv = dv[..., : dout.shape[-1]]
+        return (
+            dq,
+            dk,
+            dv,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+def flash_attn_qkvpacked_func(
+    qkv,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # <=0.0 means deactivate
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If Q, K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of Q, K, V.
+    For multi-query and grouped-query attention (MQA/GQA), please see
+    flash_attn_kvpacked_func and flash_attn_func.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.
+    Arguments:
+        qkv: (batch_size, seqlen, 3, nheads, headdim)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to
+            the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnQKVPackedFunc.apply(
+        qkv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_kvpacked_func(
+    q,
+    kv,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of K, V.
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        kv: (batch_size, seqlen, 2, nheads_k, headdim)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnKVPackedFunc.apply(
+        q,
+        kv,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_func(
+    q,
+    k,
+    v,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        k: (batch_size, seqlen, nheads_k, headdim)
+        v: (batch_size, seqlen, nheads_k, headdim)
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnFunc.apply(
+        q,
+        k,
+        v,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_varlen_qkvpacked_func(
+    qkv,
+    cu_seqlens,
+    max_seqlen,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If Q, K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_varlen_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of Q, K, V.
+    For multi-query and grouped-query attention (MQA/GQA), please see
+    flash_attn_varlen_kvpacked_func and flash_attn_varlen_func.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.
+    Arguments:
+        qkv: (total, 3, nheads, headdim), where total = total number of tokens in the batch.
+        cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into qkv.
+        max_seqlen: int. Maximum sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnVarlenQKVPackedFunc.apply(
+        qkv,
+        cu_seqlens,
+        max_seqlen,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_varlen_kvpacked_func(
+    q,
+    kv,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    If K, V are already stacked into 1 tensor, this function will be faster than
+    calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
+    of the gradients of K, V.
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.
+        kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into q.
+        cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into kv.
+        max_seqlen_q: int. Maximum query sequence length in the batch.
+        max_seqlen_k: int. Maximum key sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnVarlenKVPackedFunc.apply(
+        q,
+        kv,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        torch.is_grad_enabled(),
+    )
+def flash_attn_varlen_func(
+    q,
+    k,
+    v,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=False,
+    block_table=None,
+):
+    """dropout_p should be set to 0.0 during evaluation
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Arguments:
+        q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.
+        k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
+        cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into q.
+        cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+           of the sequences in the batch, used to index into kv.
+        max_seqlen_q: int. Maximum query sequence length in the batch.
+        max_seqlen_k: int. Maximum key sequence length in the batch.
+        dropout_p: float. Dropout probability.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        deterministic: bool. Whether to use the deterministic implementation of the backward pass,
+            which is slightly slower and uses more memory. The forward pass is always deterministic.
+        return_attn_probs: bool. Whether to return the attention probabilities. This option is for
+           testing only. The returned probabilities are not guaranteed to be correct
+           (they might not have the right scaling).
+    Return:
+        out: (total, nheads, headdim).
+        softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+        S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
+            The output of softmax (possibly with different scaling). It also encodes the dropout
+            pattern (negative means that location was dropped, nonnegative means it was kept).
+    """
+    return FlashAttnVarlenFunc.apply(
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        causal,
+        window_size,
+        softcap,
+        alibi_slopes,
+        deterministic,
+        return_attn_probs,
+        block_table,
+        False if q.device.type == "cpu" else torch.is_grad_enabled(),
+    )
+def flash_attn_with_kvcache(
+    q,
+    k_cache,
+    v_cache,
+    k=None,
+    v=None,
+    rotary_cos=None,
+    rotary_sin=None,
+    cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None,
+    cache_batch_idx: Optional[torch.Tensor] = None,
+    cache_leftpad: Optional[torch.Tensor] = None,
+    block_table: Optional[torch.Tensor] = None,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),  # -1 means infinite context window
+    softcap=0.0,  # 0.0 means deactivated
+    rotary_interleaved=True,
+    alibi_slopes=None,
+    num_splits=0,
+    return_softmax_lse=False,
+):
+    """
+    If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from
+    k and v. This is useful for incremental decoding: you can pass in the cached keys/values from
+    the previous step, and update them with the new keys/values from the current step, and do
+    attention with the updated cache, all in 1 kernel.
+    If you pass in k / v, you must make sure that the cache is large enough to hold the new values.
+    For example, the KV cache could be pre-allocated with the max sequence length, and you can use
+    cache_seqlens to keep track of the current sequence lengths of each sequence in the batch.
+    Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be
+    rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
+    If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos
+    and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
+    If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at
+    indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens).
+    See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function.
+    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
+    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
+    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
+    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
+    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
+    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
+        1 1 1 1 0
+        1 1 1 1 1
+    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
+        0 0
+        0 0
+        0 0
+        1 0
+        1 1
+    If the row of the mask is all zero, the output will be zero.
+    If window_size != (-1, -1), implements sliding window local attention. Query at position i
+    will only attend to keys between
+    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
+    Note: Does not support backward pass.
+    Arguments:
+        q: (batch_size, seqlen, nheads, headdim)
+        k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,
+            or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)
+            page_block_size must be a multiple of 256.
+        v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,
+            or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)
+        k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate
+            k with k_cache, starting at the indices specified by cache_seqlens.
+        v [optional]: (batch_size, seqlen_new, nheads_k, headdim). Similar to k.
+        rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding
+            to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16.
+        rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos.
+        cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the
+            KV cache.
+        cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache.
+            If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1].
+            If the indices are not distinct, and k and v are provided, the values updated in the cache
+                 might come from any of the duplicate indices.
+        cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0.
+        block_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32.
+        softmax_scale: float. The scaling of QK^T before applying softmax.
+            Default to 1 / sqrt(headdim).
+        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
+        softcap: float. Anything > 0 activates softcapping attention.
+        rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in.
+            If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False,
+            rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1
+            (i.e. GPT-NeoX style).
+        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
+            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
+            is added to the attention score of query i and key j.
+        num_splits: int. If > 1, split the key/value into this many chunks along the sequence.
+           If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic
+           to automatically determine the number of splits.
+           Don't change this unless you know what you are doing.
+        return_softmax_lse: bool. Whether to return the logsumexp of the attention scores.
+    Return:
+        out: (batch_size, seqlen, nheads, headdim).
+        softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The
+            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
+            normalization factor).
+    """
+    assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension"
+    assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension"
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** (-0.5)
+    if cache_seqlens is not None and isinstance(cache_seqlens, int):
+        cache_seqlens = torch.full(
+            (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device
+        )
+        cache_seqlens = maybe_contiguous(cache_seqlens)
+    cache_batch_idx = maybe_contiguous(cache_batch_idx)
+    block_table = maybe_contiguous(block_table)
+    out, softmax_lse = flash_attn.fwd_kvcache(
+        q,
+        k_cache,
+        v_cache,
+        k,
+        v,
+        cache_seqlens,
+        rotary_cos,
+        rotary_sin,
+        cache_batch_idx,
+        cache_leftpad,
+        block_table,
+        alibi_slopes,
+        None,
+        softmax_scale,
+        causal,
+        window_size[0],
+        window_size[1],
+        softcap,
+        rotary_interleaved,
+        num_splits,
+    )
+    return (out, softmax_lse) if return_softmax_lse else out

build/torch-stable-abi210-cu134-aarch64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,39 @@

+{
+  "name": "flash-attn2",
+  "id": "_flash_attn2_cuda_460047c",
+  "version": 4,
+  "kernels-minver": "0.14.0",
+  "license": "BSD-3-Clause",
+  "python-depends": [],
+  "kernel-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Nz8hqrusjoUHQqszpxfF2USMwUXn1VfpvcmtAAvutSM=",
+      "_flash_attn2_cuda_460047c.abi3.so": "kjvVCXGj2U9FuaX6vCIjp/3KHUDYzP9ISJB+uFDpm0w=",
+      "_ops.py": "DV4Eo4+s3aaNmDsfXoZiYKNLnRUJqlsdDGlMsdKFgk4=",
+      "flash_attn_interface.py": "6UYUhfUax4+yOUygyS7K8dVjrIjU5GQr1AJ2gSp8rPo=",
+      "symbols.json": "uvC+87tHUhtHv7L/Nl53T/Ms7ceX317/jwQBoqHY14A="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.18.0-dev0",
+      "commit": "d054906a13710b811b96697b4635f03020842733",
+      "dirty": false
+    },
+    "kernel": {
+      "commit": "460047c291fc7f0b3dbd4239b46bd5e1281b1806",
+      "dirty": false
+    }
+  }
+}

build/torch-stable-abi210-cu134-aarch64-linux/metadata.json.sigstore ADDED Viewed

	@@ -0,0 +1 @@

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHLDCCBrKgAwIBAgIUfIlqLNu6UdXYNtVUv5fTcZRypY4wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDI0MDAwNzUzWhcNMjYxMDI0MDAxNzUzWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEpcQphv3Uy88KqfPC2bzQt1HLnip0QId3cZg/mMIXU35xCKnr7Z3cdW02SOZ8g8HSD5YMT9K1bDeEeJGBzXagCqOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBCg0NjAwNDdjMjkxZmM3ZjBiM2RiZDQyMzliNDZiZDVlMTI4MWIxODA2MBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDCg0NjAwNDdjMjkxZmM3ZjBiM2RiZDQyMzliNDZiZDVlMTI4MWIxODA2MBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDCg0NjAwNDdjMjkxZmM3ZjBiM2RiZDQyMzliNDZiZDVlMTI4MWIxODA2MB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoNDYwMDQ3YzI5MWZjN2YwYjNkYmQ0MjM5YjQ2YmQ1ZTEyODFiMTgwNjAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTk4MjY0MDUxNS9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhazxDoAAABAMARzBFAiAwLD/CeVrpcA2C2j58s920o2EJYCRTjPvcMRrbsTS6RwIhAPfV39SmThkryIVjNRuewtpScIsCH/2NACYy+FMTtc+GMAoGCCqGSM49BAMDA2gAMGUCMGv48brzSXkUdf5w2ftOJMGQQiZhmnV1/BvuBvoSNKxmjnCwI50V0ijkF5zcDNcD/QIxANLbDO4W25uC2/31J9S8JTlIx+8apDRx5Gide5gHyzLSgV1lj4/1jM5pN2u+aoI2Gg=="},"tlogEntries":[{"logIndex":"3430233537","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792800474","inclusionPromise":{"signedEntryTimestamp":"MEQCIGf4qRksR2TOvKJqv5C3OXL9KCWq4DohDNE5OiRqroDSAiBZ/na5GYKHZMVmMM/RBb1tA3aWMI0vOAvXzqPbefcYcQ=="},"inclusionProof":{"logIndex":"3308329275","rootHash":"gvJkaoo5K4nnBj5tgxAdNtAsuOVBE3dwxp4HnZkfvek=","treeSize":"3308329751","hashes":["A6lD2MPXJCkNJ03unZasE28kCBO4cEhrMfr35263otA=","2juSPfQ+oiHAxdHnLkvjhIT+ZK12WdDckwwg3N92P7A=","LifhnEL2YNR6i6SL1PtcvHuNSG8cMuLHlq/xkxrLIl8=","bsl+fw0R4ouIsjrASTpknFPjkE8rV44VNY9PzMjKZXY=","7m7Os3c6xSIltD1PYdtmZWTvYcp/y98WcmkA5VHZLMc=","NrwiN5u5p9JLACtCux17Ga5le7OR7KG91Swm1vKzN4I=","B+n7TQ5+p34ZRpHoH7r7iRTI8eA90txEiNvVk/XQGlo=","ornPIqk+dhfmtF5rp9z75eQBaUggI9Zd0Yoov3BG4Zw=","0YVZixFviDG94SVGteh82jVzdNtBsMQZbb6BiHiUA30=","2Vc45+NdnNK2xlrvQ9A2rn3JmookZpsttqolVggzpFM=","eIbVG1biupwReGylDyasAnDJDAjvcA3MN0/o/j5H/Vg=","gdS39fsn5MOVlKluUgGY2sdKF3H5oauxkHMa7DSSNcg=","8x0N/oAmvpY+6DQUGTU+i6a1Cuqphro2WbV5WZ/RHL8=","DMchGrSAeKV9qHo75XJV3UFZ/CHF8zofPRu+NgVRwA4=","V6dXjeDcwYxBG/9gVSBM8rYdPrAqGpZsf1e7v6l67GA=","iDfImp65w8dirRfYGNgFA34JTUYDWbcJwNBBU/Ok+X4=","YYCr61A8azar3dfO8Ct+pri7LqDxkUlUkfs7fiqYErU=","1DaDKzrCNNBhmSdY5vFjz+NpAmxmYfYEgBCdPLrMP1A=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3308329751\ngvJkaoo5K4nnBj5tgxAdNtAsuOVBE3dwxp4HnZkfvek=\n\n— rekor.sigstore.dev wNI9ajBEAiA7HK7JgdGaQLRpTJ61zoG8UJzLcQbCNxmOd2H6oBto3QIgPoF4fB3lpbTOp7T38nRRszHe0SI4Tet8YnGxCwxSXWI=\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIyNTI4NjhhNGQyOGI4NzZlYzYyZjMxNmUzMWRkYTA5ZDQ5ZTlmNDYxYWQwNTk1MzlmNTc2YjJkYjJlZTE1YWMxIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUURsaktTUFVEQWtlYU9TMzFTTlFSMWt6amhrMGpwaVgvQy95bENaRS9zVW5nSWhBTHo4SEFwalYyM0E5ei9uRHZZNFRGdWZ6K2YvalBobkZJZkozNnRaRS9TTiIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1SRU5EUW5KTFowRjNTVUpCWjBsVlprbHNjVXhPZFRaVlpGaFpUblJXVlhZMVpsUmpXbEo1Y0ZrMGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFU1RCTlJFRjNUbnBWZWxkb1kwNU5hbGw0VFVSSk1FMUVRWGhPZWxWNlYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ3WTFGd2FIWXpWWGs0T0V0eFpsQkRNbUo2VVhReFNFeHVhWEF3VVVsa00yTmFaeThLYlUxSldGVXpOWGhEUzI1eU4xb3pZMlJYTURKVFQxbzRaemhJVTBRMVdVMVVPVXN4WWtSbFJXVktSMEo2V0dGblEzRlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmN3VG1wQmR3cE9SR1JxVFdwcmVGcHRUVE5hYWtKcFRUSlNhVnBFVVhsTmVteHBUa1JhYVZwRVZteE5WRWswVFZkSmVFOUVRVEpOUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEWnpCT2FrRjNUa1JrYWsxcWEzaGFiVTB6Q2xwcVFtbE5NbEpwV2tSUmVVMTZiR2xPUkZwcFdrUldiRTFVU1RSTlYwbDRUMFJCTWsxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEWnpCT2FrRjNDazVFWkdwTmFtdDRXbTFOTTFwcVFtbE5NbEpwV2tSUmVVMTZiR2xPUkZwcFdrUldiRTFVU1RSTlYwbDRUMFJCTWsxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlPUkZsM1RVUlJNMWw2U1RWTlYxcHFUakpaZDFscVRtdFpiVkV3VFdwTk5WbHFVVEpaYlZFeFdsUkZlVTlFUm1rS1RWUm5kMDVxUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUnJORTFxV1RCTlJGVjRUbE01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvWVhwNFJHOUJRVUZDUVUxQlVucENSa0ZwUVhkTVJDOURaVlp5Y0dOQk1rTXlhalU0Y3preU1HOHlSVXBaUTFKVWFsQjJZd3BOVW5KaWMxUlRObEozU1doQlVHWldNemxUYlZSb2EzSjVTVlpxVGxKMVpYZDBjRk5qU1hORFNDOHlUa0ZEV1hrclJrMVVkR01yUjAxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKblFVMUhWVU5OUjNZME9HSnllbE5ZYTFWa1pqVjNNbVowVDBwTlIxRlJhVnBvYlc1V01TOUNkblZDZG05VFRrdDRiV3B1UTNjS1NUVXdWakJwYW10R05YcGpSRTVqUkM5UlNYaEJUa3hpUkU4MFZ6STFkVU15THpNeFNqbFRPRXBVYkVsNEt6aGhjRVJTZURWSGFXUmxOV2RJZVhwTVV3cG5WakZzYWpRdk1XcE5OWEJPTW5VcllXOUpNa2RuUFQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgxAQshCKRlI1Pb2bAj9KFX31k9RqTG+GzEdk6zzSe9FQCFCRfS1fC3kY93BLQYuBiV9kW0SCsGA8yMDI2MTAyNDAwMDc1NFowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDI0MDAwNzU0WjAvBgkqhkiG9w0BCQQxIgQgKYUPinZ7u9mF7M+Vqfl8o6ym/pWAGm/AA/NVEIq4DiowgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxAOfOfD0hEGMls7xwJHci7OLQt6Qvku2I1Nb2833ad+2sRiZplkmYRU975nVt0EqtBwIxAInS0x2jMo+2BSef7XYiRS+na+f//QqAjBUE3FBJqXjdgCv8yOVhN0+f26chU1s4og=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"JShopNKLh27GLzFuMd2gnUnp9GGtBZU59Xay2y7hWsE="},"signature":"MEYCIQDljKSPUDAkeaOS31SNQR1kzjhk0jpiX/C/ylCZE/sUngIhALz8HApjV23A9z/nDvY4TFufz+f/jPhnFIfJ36tZE/SN"}}

build/torch-stable-abi210-cu134-aarch64-linux/symbols.json ADDED Viewed

	@@ -0,0 +1,1273 @@

+{
+  "schema_version": 1,
+  "module": "_flash_attn2_cuda_460047c",
+  "functions": [
+    {
+      "name": "flash_attn_func",
+      "qualname": "flash_attn_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (batch_size, seqlen, nheads, headdim)\n    k: (batch_size, seqlen, nheads_k, headdim)\n    v: (batch_size, seqlen, nheads_k, headdim)\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_kvpacked_func",
+      "qualname": "flash_attn_kvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of K, V.\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (batch_size, seqlen, nheads, headdim)\n    kv: (batch_size, seqlen, 2, nheads_k, headdim)\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "kv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_qkvpacked_func",
+      "qualname": "flash_attn_qkvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf Q, K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of Q, K, V.\nFor multi-query and grouped-query attention (MQA/GQA), please see\nflash_attn_kvpacked_func and flash_attn_func.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.\n\nArguments:\n    qkv: (batch_size, seqlen, 3, nheads, headdim)\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to\n        the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "qkv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_varlen_func",
+      "qualname": "flash_attn_varlen_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.\n    k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.\n    v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.\n    cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into q.\n    cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into kv.\n    max_seqlen_q: int. Maximum query sequence length in the batch.\n    max_seqlen_k: int. Maximum key sequence length in the batch.\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (total, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_varlen_kvpacked_func",
+      "qualname": "flash_attn_varlen_kvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of K, V.\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n    q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.\n    kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch.\n    cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into q.\n    cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into kv.\n    max_seqlen_q: int. Maximum query sequence length in the batch.\n    max_seqlen_k: int. Maximum key sequence length in the batch.\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (total, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "kv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_varlen_qkvpacked_func",
+      "qualname": "flash_attn_varlen_qkvpacked_func",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "dropout_p should be set to 0.0 during evaluation\nIf Q, K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_varlen_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of Q, K, V.\nFor multi-query and grouped-query attention (MQA/GQA), please see\nflash_attn_varlen_kvpacked_func and flash_attn_varlen_func.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.\n\nArguments:\n    qkv: (total, 3, nheads, headdim), where total = total number of tokens in the batch.\n    cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths\n       of the sequences in the batch, used to index into qkv.\n    max_seqlen: int. Maximum sequence length in the batch.\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|)\n        is added to the attention score of query i and key j.\n    deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n        which is slightly slower and uses more memory. The forward pass is always deterministic.\n    return_attn_probs: bool. Whether to return the attention probabilities. This option is for\n       testing only. The returned probabilities are not guaranteed to be correct\n       (they might not have the right scaling).\nReturn:\n    out: (total, nheads, headdim).\n    softmax_lse [optional, if return_attn_probs=True]: (nheads, total_q_seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).\n    S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).\n        The output of softmax (possibly with different scaling). It also encodes the dropout\n        pattern (negative means that location was dropped, nonnegative means it was kept).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "qkv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "cu_seqlens",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "max_seqlen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "dropout_p",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "return_attn_probs",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "flash_attn_with_kvcache",
+      "qualname": "flash_attn_with_kvcache",
+      "module": "_flash_attn2_cuda_460047c.flash_attn_interface",
+      "kind": "function",
+      "doc": "If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from\nk and v. This is useful for incremental decoding: you can pass in the cached keys/values from\nthe previous step, and update them with the new keys/values from the current step, and do\nattention with the updated cache, all in 1 kernel.\n\nIf you pass in k / v, you must make sure that the cache is large enough to hold the new values.\nFor example, the KV cache could be pre-allocated with the max sequence length, and you can use\ncache_seqlens to keep track of the current sequence lengths of each sequence in the batch.\n\nAlso apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be\nrotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.\nIf causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos\nand rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.\nIf not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at\nindices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens).\n\nSee tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function.\n\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nIf causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n    1 1 1 1 0\n    1 1 1 1 1\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n    0 0\n    0 0\n    0 0\n    1 0\n    1 1\nIf the row of the mask is all zero, the output will be zero.\n\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nNote: Does not support backward pass.\n\nArguments:\n    q: (batch_size, seqlen, nheads, headdim)\n    k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,\n        or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)\n        page_block_size must be a multiple of 256.\n    v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,\n        or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)\n    k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate\n        k with k_cache, starting at the indices specified by cache_seqlens.\n    v [optional]: (batch_size, seqlen_new, nheads_k, headdim). Similar to k.\n    rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding\n        to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16.\n    rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos.\n    cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the\n        KV cache.\n    cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache.\n        If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1].\n        If the indices are not distinct, and k and v are provided, the values updated in the cache\n             might come from any of the duplicate indices.\n    cache_leftpad: (batch_size,), dtype torch.int32. The index that the KV cache starts. If None, assume 0.\n    block_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 / sqrt(headdim).\n    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n    window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n    softcap: float. Anything > 0 activates softcapping attention.\n    rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in.\n        If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False,\n        rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1\n        (i.e. GPT-NeoX style).\n    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n        is added to the attention score of query i and key j.\n    num_splits: int. If > 1, split the key/value into this many chunks along the sequence.\n       If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic\n       to automatically determine the number of splits.\n       Don't change this unless you know what you are doing.\n    return_softmax_lse: bool. Whether to return the logsumexp of the attention scores.\n\nReturn:\n    out: (batch_size, seqlen, nheads, headdim).\n    softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The\n        logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax\n        normalization factor).",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k_cache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "v_cache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": null
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "rotary_cos",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "rotary_sin",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "cache_seqlens",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "int | torch.Tensor | None"
+          },
+          {
+            "name": "cache_batch_idx",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "cache_leftpad",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          },
+          {
+            "name": "window_size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "(-1, -1)",
+            "annotation": null
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": null
+          },
+          {
+            "name": "rotary_interleaved",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "True",
+            "annotation": null
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": null
+          },
+          {
+            "name": "num_splits",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": null
+          },
+          {
+            "name": "return_softmax_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": null
+          }
+        ],
+        "return_annotation": null
+      }
+    },
+    {
+      "name": "fwd",
+      "qualname": "fwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Forward pass for multi-head attention.\n\nArgs:\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    out: Optional output tensor, same shape as q\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    return_softmax: Whether to return softmax weights\n    gen: Optional random number generator\n\nReturns:\n    List of tensors: [output, softmax_lse, (softmax if return_softmax)]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "return_softmax",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "varlen_fwd",
+      "qualname": "varlen_fwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Forward pass for multi-head attention with variable sequence lengths.\n\nArgs:\n    q: Query tensor of shape [total_q, num_heads, head_size]\n    k: Key tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    v: Value tensor of shape [total_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]\n    cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]\n    out: Optional output tensor of shape [total_q, num_heads, head_size]\n    seqused_k: Optional tensor specifying how many keys to use per batch element [batch_size]\n    leftpad_k: Optional left padding for keys of shape [batch_size]\n    block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    max_seqlen_q: Maximum sequence length for queries\n    max_seqlen_k: Maximum sequence length for keys\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    zero_tensors: Whether to zero tensors before computation\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    return_softmax: Whether to return softmax weights\n    gen: Optional random number generator\n\nReturns:\n    List of tensors: [output, softmax_lse, (softmax if return_softmax)]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "seqused_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "leftpad_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "zero_tensors",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "return_softmax",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "bwd",
+      "qualname": "bwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Backward pass for multi-head attention.\n\nArgs:\n    dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]\n    softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]\n    dq: Optional gradient tensor for queries, same shape as q\n    dk: Optional gradient tensor for keys, same shape as k\n    dv: Optional gradient tensor for values, same shape as v\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    deterministic: Whether to use deterministic algorithms\n    gen: Optional random number generator\n    rng_state: Optional RNG state from forward pass\n\nReturns:\n    List of tensors: [dq, dk, dv]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "dout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "softmax_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "dq",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dk",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          },
+          {
+            "name": "rng_state",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "varlen_bwd",
+      "qualname": "varlen_bwd",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Backward pass for multi-head attention with variable sequence lengths.\n\nArgs:\n    dout: Gradient tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    k: Key tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    v: Value tensor of shape [batch_size, seqlen_k, num_heads_k, head_size]\n    out: Output tensor from forward pass of shape [batch_size, seqlen_q, num_heads, head_size]\n    softmax_lse: Log-sum-exp values from forward pass of shape [batch_size, num_heads, seqlen_q]\n    cu_seqlens_q: Cumulative sequence lengths for queries of shape [batch_size+1]\n    cu_seqlens_k: Cumulative sequence lengths for keys of shape [batch_size+1]\n    dq: Optional gradient tensor for queries, same shape as q\n    dk: Optional gradient tensor for keys, same shape as k\n    dv: Optional gradient tensor for values, same shape as v\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    max_seqlen_q: Maximum sequence length for queries\n    max_seqlen_k: Maximum sequence length for keys\n    p_dropout: Dropout probability\n    softmax_scale: Scale factor for softmax\n    zero_tensors: Whether to zero tensors before computation\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    deterministic: Whether to use deterministic algorithms\n    gen: Optional random number generator\n    rng_state: Optional RNG state from forward pass\n\nReturns:\n    List of tensors: [dq, dk, dv]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "dout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "softmax_lse",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "cu_seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "dq",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dk",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "dv",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "max_seqlen_q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "max_seqlen_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0",
+            "annotation": "int"
+          },
+          {
+            "name": "p_dropout",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "zero_tensors",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "deterministic",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "gen",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch._C.Generator | None"
+          },
+          {
+            "name": "rng_state",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    },
+    {
+      "name": "fwd_kvcache",
+      "qualname": "fwd_kvcache",
+      "module": "_flash_attn2_cuda_460047c",
+      "kind": "function",
+      "doc": "Forward pass for multi-head attention with KV cache.\n\nArgs:\n    q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size]\n    kcache: Key cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    vcache: Value cache tensor of shape [batch_size_c, seqlen_k, num_heads_k, head_size] or [num_blocks, page_block_size, num_heads_k, head_size]\n    k: Optional new keys tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]\n    v: Optional new values tensor of shape [batch_size, seqlen_knew, num_heads_k, head_size]\n    seqlens_k: Optional sequence lengths for keys of shape [batch_size]\n    rotary_cos: Optional rotary cosine tensor of shape [seqlen_ro, rotary_dim/2]\n    rotary_sin: Optional rotary sine tensor of shape [seqlen_ro, rotary_dim/2]\n    cache_batch_idx: Optional indices to index into the KV cache\n    leftpad_k: Optional left padding for keys of shape [batch_size]\n    block_table: Optional block table of shape [batch_size, max_num_blocks_per_seq]\n    alibi_slopes: Optional ALiBi slopes tensor of shape [num_heads] or [batch_size, num_heads]\n    out: Optional output tensor, same shape as q\n    softmax_scale: Scale factor for softmax\n    is_causal: Whether to use causal attention\n    window_size_left: Window size for left context (-1 for unlimited)\n    window_size_right: Window size for right context (-1 for unlimited)\n    softcap: Soft cap for attention weights\n    is_rotary_interleaved: Whether rotary embeddings are interleaved\n    num_splits: Number of splits for computation\n\nReturns:\n    List of tensors: [output, softmax_lse]",
+      "signature": {
+        "parameters": [
+          {
+            "name": "q",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "kcache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "vcache",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "v",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "seqlens_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "rotary_cos",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "rotary_sin",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "cache_batch_idx",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "leftpad_k",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "block_table",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "alibi_slopes",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "out",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "softmax_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "float | None"
+          },
+          {
+            "name": "is_causal",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "window_size_left",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "window_size_right",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "-1",
+            "annotation": "int"
+          },
+          {
+            "name": "softcap",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "0.0",
+            "annotation": "float"
+          },
+          {
+            "name": "is_rotary_interleaved",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "False",
+            "annotation": "bool"
+          },
+          {
+            "name": "num_splits",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "1",
+            "annotation": "int"
+          }
+        ],
+        "return_annotation": "List[torch.Tensor]"
+      }
+    }
+  ],
+  "layers": []
+}