JingzeShi commited on Jun 21

Commit

a61996a

verified ·

1 Parent(s): 28eccca

Uploaded using `kernel-builder`.

Browse files

Files changed (33) hide show

build/torch-cuda/__init__.py +29 -0
build/torch-cuda/_ops.py +38 -0
build/torch-cuda/activations.py +296 -0
build/torch-cuda/assert_inputs.py +327 -0
build/torch-cuda/autotuner.py +360 -0
build/torch-cuda/block_info.py +277 -0
build/torch-cuda/cache_utils.py +180 -0
build/torch-cuda/flash_bwd_postprocess.py +215 -0
build/torch-cuda/flash_bwd_preprocess.py +266 -0
build/torch-cuda/flash_dec_combine.py +222 -0
build/torch-cuda/flash_dense_bwd.py +1019 -0
build/torch-cuda/flash_dense_dec.py +857 -0
build/torch-cuda/flash_dense_fwd.py +957 -0
build/torch-cuda/flash_fwd_combine.py +239 -0
build/torch-cuda/flash_gated_bwd.py +1317 -0
build/torch-cuda/flash_gated_dec.py +1218 -0
build/torch-cuda/flash_gated_fwd.py +1337 -0
build/torch-cuda/flash_proj_bwd.py +21 -0
build/torch-cuda/flash_proj_fwd.py +45 -0
build/torch-cuda/flash_sparse_attention/__init__.py +26 -0
build/torch-cuda/flash_sparse_bwd.py +1041 -0
build/torch-cuda/flash_sparse_dec.py +871 -0
build/torch-cuda/flash_sparse_fwd.py +974 -0
build/torch-cuda/interface.py +1912 -0
build/torch-cuda/kernel_repr.py +115 -0
build/torch-cuda/launch_grid.py +202 -0
build/torch-cuda/launch_template.py +457 -0
build/torch-cuda/mask.py +106 -0
build/torch-cuda/metadata.json +47 -0
build/torch-cuda/quant.py +66 -0
build/torch-cuda/scheduler.py +2609 -0
build/torch-cuda/seqlen_info.py +273 -0
build/torch-cuda/utils.py +167 -0

build/torch-cuda/__init__.py ADDED Viewed

	@@ -0,0 +1,29 @@

+from .interface import (
+    flash_dense_attn_func,
+    flash_dense_attn_with_kvcache_func,
+    flash_dense_attn_varlen_func,
+    flash_dense_attn_varlen_with_kvcache_func,
+    flash_sparse_attn_func,
+    flash_sparse_attn_with_kvcache_func,
+    flash_sparse_attn_varlen_func,
+    flash_sparse_attn_varlen_with_kvcache_func,
+    flash_gated_attn_func,
+    flash_gated_attn_with_kvcache_func,
+    flash_gated_attn_varlen_func,
+    flash_gated_attn_varlen_with_kvcache_func,
+)
+__all__ = [
+    "flash_dense_attn_func",
+    "flash_dense_attn_with_kvcache_func",
+    "flash_dense_attn_varlen_func",
+    "flash_dense_attn_varlen_with_kvcache_func",
+    "flash_sparse_attn_func",
+    "flash_sparse_attn_with_kvcache_func",
+    "flash_sparse_attn_varlen_func",
+    "flash_sparse_attn_varlen_with_kvcache_func",
+    "flash_gated_attn_func",
+    "flash_gated_attn_with_kvcache_func",
+    "flash_gated_attn_varlen_func",
+    "flash_gated_attn_varlen_with_kvcache_func",
+]

build/torch-cuda/_ops.py ADDED Viewed

	@@ -0,0 +1,38 @@

+import torch
+def get_backend() -> str:
+    """Detect the backend by inspecting torch."""
+    import torch
+    if hasattr(torch, "neuron"):
+        # Needs to be sorted before specific Torch builds, since Neuron
+        # extension can be loaded into e.g. CUDA Torch builds.
+        return "neuron"
+    elif torch.version.cuda is not None:
+        return "cuda"
+    elif torch.version.hip is not None:
+        return "rocm"
+    elif torch.backends.mps.is_available():
+        return "metal"
+    elif hasattr(torch.version, "xpu") and torch.version.xpu is not None:
+        return "xpu"
+    else:
+        return "cpu"
+def _find_ops_name() -> str:
+    kernel_name = "flash_sparse_attention"
+    unique_id = "7db628a_dirty"
+    backend = get_backend()
+    return f"_{kernel_name}_{backend}_{unique_id}"
+_OPS_NAME = _find_ops_name()
+ops = getattr(torch.ops, _OPS_NAME)
+def add_op_namespace_prefix(op_name: str) -> str:
+    """
+    Prefix op by namespace.
+    """
+    return f"{_OPS_NAME}::{op_name}"

build/torch-cuda/activations.py ADDED Viewed

	@@ -0,0 +1,296 @@

+import triton
+import triton.language as tl
+@triton.jit
+def check_inf(x):
+    return tl.maximum(x, -1e6)
+@triton.jit
+def exp2(x):
+    """
+    Compute 2^x, using fast hardware approximation on CUDA and default implementation on other backends.
+    :param x: Input tensor of shape [BLOCK_M, BLOCK_N].
+    :return: Tensor of shape [BLOCK_M, BLOCK_N] containing 2^x values.
+    """
+    if tl.target_info.is_cuda():
+        return tl.inline_asm_elementwise(
+            "ex2.approx.ftz.f32 $0, $1;",
+            "=r, r",
+            [x],
+            dtype=tl.float32,
+            is_pure=True,
+            pack=1,
+        )
+    else:
+        return tl.exp2(x)
+@triton.jit
+def exp(x):
+    """
+    Compute e^x, using fast hardware approximation on CUDA and default implementation on other backends.
+    :param x: Input tensor of shape [BLOCK_M, BLOCK_N].
+    :return: Tensor of shape [BLOCK_M, BLOCK_N] containing e^x values.
+    """
+    log2_e: tl.constexpr = 1.4426950408889634
+    x *= log2_e
+    return exp2(x)
+@triton.jit
+def online_softmax(
+    acc_s,
+    row_max,
+    row_sum,
+    scale_log2,
+    CHECK_INF: tl.constexpr,
+    RESCALE_THRESHOLD: tl.constexpr,
+):
+    """
+    Apply online softmax to acc_s, and update row_max and row_sum.
+    :param acc_s: Attention scores tensor of shape [BLOCK_M, BLOCK_N].
+    :param row_max: Current maximum values per row of shape [BLOCK_M], init to -inf.
+    :param row_sum: Current sum values per row of shape [BLOCK_M], init to 0.
+    :param scale_log2: Log2 of the scaling factor to be applied to acc_s.
+    :param CHECK_INF: Boolean flag indicating if -inf row_max should be clamped to 0.
+    :param RESCALE_THRESHOLD: Threshold for rescaling to avoid underflow. If <= 0, rescaling is disabled.
+    :return p: Softmax probabilities tensor of shape [BLOCK_M, BLOCK_N].
+    :return row_max_new: Updated maximum values per row of shape [BLOCK_M].
+    :return row_sum_new: Updated sum values per row of shape [BLOCK_M].
+    :return row_scale: Scaling factors per row of shape [BLOCK_M].
+    """
+    # Compute current row max
+    row_max_curr = tl.max(acc_s, axis=1)
+    # Update row max
+    row_max_new = tl.maximum(row_max_curr, row_max)
+    # Avoid exp(-inf - (-inf)) = nan by clamping -inf to 0
+    if CHECK_INF:
+        row_max_new = check_inf(row_max_new)
+    # Compute scaled differences to new row max
+    acc_scale_log2 = (row_max - row_max_new) * scale_log2
+    # Compute row scale
+    if RESCALE_THRESHOLD > 0.0:
+        # Triton can only skip computation at block granularity
+        if tl.min(acc_scale_log2) < -RESCALE_THRESHOLD:
+            row_scale = exp2(acc_scale_log2)
+        else:
+            row_max_new = row_max
+            row_scale = acc_scale_log2 * 0.0 + 1.0
+    else:
+        row_scale = exp2(acc_scale_log2)
+    # Compute attention weights
+    p = exp2(acc_s * scale_log2 - row_max_new[:, None] * scale_log2)
+    # Update row sum
+    row_sum_cur = tl.sum(p, axis=1)
+    row_sum_new = row_sum * row_scale + row_sum_cur
+    return p, row_max_new, row_sum_new, row_scale
+@triton.jit
+def online_sparse_softmax(
+    acc_s,
+    row_max,
+    row_sum,
+    scale_log2,
+    softmax_threshold_log2,
+    CHECK_INF: tl.constexpr,
+):
+    """
+    Apply online sparse softmax to acc_s, and update row_max and row_sum.
+    :param acc_s: Attention scores tensor of shape [BLOCK_M, BLOCK_N].
+    :param row_max: Current maximum values per row of shape [BLOCK_M], init to -inf.
+    :param row_sum: Current sum values per row of shape [BLOCK_M], init to 0.
+    :param scale_log2: Log2 of the scaling factor to be applied to acc_s.
+    :param softmax_threshold_log2: Threshold in log2-domain for block-level skip.
+    :param CHECK_INF: Boolean flag indicating if -inf row_max should be clamped to 0.
+    :return p: Softmax probabilities tensor of shape [BLOCK_M, BLOCK_N].
+    :return row_max_new: Updated maximum values per row of shape [BLOCK_M].
+    :return row_sum_new: Updated sum values per row of shape [BLOCK_M].
+    :return row_scale: Scaling factors per row of shape [BLOCK_M].
+    :return skip_softmax: Boolean indicating whether this block was skipped.
+    """
+    # Compute current row max
+    row_max_curr = tl.max(acc_s, axis=1)
+    # Update skip condition based on threshold
+    row_max_diff_log2 = (row_max_curr - row_max) * scale_log2
+    skip_softmax = tl.max(row_max_diff_log2 - softmax_threshold_log2) < 0.0
+    # Return zero attention weights
+    if skip_softmax:
+        p = acc_s * 0.0
+        row_max_new = row_max
+        row_sum_new = row_sum
+        row_scale = row_max * 0.0 + 1.0
+    else:
+        # Update row max
+        row_max_new = tl.maximum(row_max_curr, row_max)
+        # Avoid exp(-inf - (-inf)) = nan by clamping -inf to 0
+        if CHECK_INF:
+            row_max_new = check_inf(row_max_new)
+        # Compute scaled differences to new row max
+        acc_scale_log2 = (row_max - row_max_new) * scale_log2
+        # Compute row scale
+        row_scale = exp2(acc_scale_log2)
+        # Compute attention weights
+        p = exp2(acc_s * scale_log2 - row_max_new[:, None] * scale_log2)
+        # Update row sum
+        row_sum_cur = tl.sum(p, axis=1)
+        row_sum_new = row_sum * row_scale + row_sum_cur
+    return p, row_max_new, row_sum_new, row_scale, skip_softmax
+@triton.jit
+def finalize(
+    row_max,
+    row_sum,
+    scale_log2,
+    final_scale,
+    IS_LOG2: tl.constexpr,
+    CHECK_NAN: tl.constexpr,
+):
+    """
+    Finalize online softmax by computing output scale and logsumexp.
+    :param row_max: Final maximum values per row of shape [BLOCK_M].
+    :param row_sum: Final sum values per row of shape [BLOCK_M].
+    :param final_scale: Scaling factor to be applied to the output.
+    :param IS_LOG2: Boolean flag indicating if the returned logsumexp should be in log2-space.
+    :param CHECK_NAN: Boolean flag indicating if nan values in row_sum should be checked and set to 1 to avoid returning nan.
+    :return row_scale: Final scaling factors per row of shape [BLOCK_M].
+    :return lse: Logsumexp values per row of shape [BLOCK_M].
+    """
+    # if row_sum is zero or nan, set it to 1 to avoid division by zero
+    if CHECK_NAN:
+        acc_o_is_zero_or_nan = (row_sum == 0.0) | (row_sum != row_sum)
+        row_scale = tl.where(acc_o_is_zero_or_nan, 1.0, final_scale / row_sum)
+        lse = tl.where(
+            acc_o_is_zero_or_nan,
+            float("-inf"),
+            row_max * scale_log2 + tl.log2(row_sum),
+        )
+    else:
+        row_scale = final_scale / row_sum
+        lse = row_max * scale_log2 + tl.log2(row_sum)
+    if not IS_LOG2:
+        # ln2 = math.log(2.0)
+        ln2: tl.constexpr = 0.6931471805599453
+        lse *= ln2
+    return row_scale, lse
+@triton.jit
+def rescale_o(
+    acc_o,
+    row_scale,
+    LAZY_RESCALE: tl.constexpr,
+):
+    """
+    Rescale output accumulator by row_scale.
+    :param acc_o: Output accumulator tensor of shape [BLOCK_M, BLOCK_N].
+    :param row_scale: Scaling factors per row of shape [BLOCK_M].
+    :param LAZY_RESCALE: Boolean flag indicating if rescaling should be applied lazily.
+    :return: Rescaled output accumulator tensor of shape [BLOCK_M, BLOCK_N].
+    """
+    if LAZY_RESCALE:
+        # On some hardware, the check can cause severe performance regression
+        if tl.min(row_scale) != 1.0:
+            acc_o = acc_o * row_scale[:, None]
+    else:
+        acc_o = acc_o * row_scale[:, None]
+    return acc_o
+@triton.jit
+def sigmoid(x):
+    """
+    Compute sigmoid of x.
+    :param x: Input tensor of shape [BLOCK_M, BLOCK_N].
+    :return: Tensor of shape [BLOCK_M, BLOCK_N] containing sigmoid values.
+    """
+    return 1.0 / (1.0 + exp(-x))
+@triton.jit
+def log_sigmoid(x, FASTMATH: tl.constexpr):
+    """
+    Compute log-sigmoid of x.
+    :param x: Input tensor of shape [BLOCK_M, BLOCK_N].
+    :param FASTMATH: Boolean flag indicating if the fast approximation should be used.
+    :return: Tensor of shape [BLOCK_M, BLOCK_N] containing log-sigmoid values.
+    """
+    if FASTMATH:
+        xc = tl.minimum(tl.abs(x), 4.0)
+        xc2 = xc * xc
+        out = tl.minimum(x, 0.0) - 0.05674870 * xc2 + 0.37664706 * xc - 0.65169323
+        return out
+    else:
+        out = tl.minimum(x, 0.0) - tl.log(1.0 + exp(-tl.abs(x)))
+        return out
+@triton.jit
+def online_gate(
+    a_max,
+    a_min,
+    d_max,
+    d_min,
+    gate_max,
+    scale_log2,
+    gate_threshold_log2,
+):
+    """
+    Determine whether to skip gate computation for the current tile based on the maximum possible gate value.
+    :param a_max: Maximum value of the Alpha tile.
+    :param a_min: Minimum value of the Alpha tile.
+    :param d_max: Maximum value of the Delta tile.
+    :param d_min: Minimum value of the Delta tile.
+    :param gate_max: Maximum value of the gate.
+    :param scale_log2: Log2 of the scaling factor applied to the gate max.
+    :param gate_threshold_log2: Threshold in log2-domain for gate-level skip.
+    :return gate_max_new: Updated gate max value after considering current tile.
+    :return skip_gate: Boolean indicating whether to skip the gate computation for this tile.
+    """
+    gate_max_curr = tl.maximum(
+        tl.maximum(a_max * d_max, a_max * d_min),
+        tl.maximum(a_min * d_max, a_min * d_min),
+    )
+    gate_max_diff_log2 = (gate_max_curr - gate_max) * scale_log2
+    skip_gate = gate_max_diff_log2 < gate_threshold_log2
+    if skip_gate:
+        gate_max_new = gate_max
+    else:
+        gate_max_new = tl.maximum(gate_max_curr, gate_max)
+    return gate_max_new, skip_gate

build/torch-cuda/assert_inputs.py ADDED Viewed

	@@ -0,0 +1,327 @@

+from typing import Optional
+import torch
+def assert_fwd_inputs(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    alpha: Optional[torch.Tensor] = None,
+    delta: Optional[torch.Tensor] = None,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    num_heads_q: int = None,
+    num_heads_kv: int = None,
+    head_dim: int = None,
+    is_quant: bool = False,
+    device: torch.device = None,
+):
+    """
+    Assert the validity of inputs for the forward kernel.
+    :param query: Query tensor
+    :param key: Key tensor
+    :param value: Value tensor
+    :param query_scale: Optional query scale tensor for quantized inputs
+    :param key_scale: Optional key scale tensor for quantized inputs
+    :param value_scale: Optional value scale tensor for quantized inputs
+    :param window_sizes: Optional window sizes tensor for local attention
+    :param alpha: Alpha tensor for gated attention
+    :param delta: Delta tensor for gated attention
+    :param cu_seqlens_q: Cumulative sequence lengths for queries
+    :param cu_seqlens_k: Cumulative sequence lengths for keys
+    :param seqused_q: Sequence used for queries
+    :param seqused_k: Sequence used for keys
+    :param num_heads_q: Number of query heads
+    :param num_heads_kv: Number of key/value heads
+    :param head_dim: Head dimension
+    :param is_quant: Whether the inputs are quantized
+    :param device: Device of the tensors
+    :raises AssertionError: If any of the assertions fail
+    """
+    assert device == query.device == key.device == value.device, (
+        "All inputs must be on the same device"
+    )
+    if is_quant:
+        assert query.dtype in [torch.float8_e5m2, torch.float8_e4m3fn], (
+            "Input dtype must be float8_e5m2 or float8_e4m3fn"
+        )
+        assert (
+            query_scale is not None
+            and key_scale is not None
+            and value_scale is not None
+        ), "Scale tensors must be provided for quantized inputs"
+        assert query_scale.dtype in [torch.float16, torch.bfloat16, torch.float32], (
+            "Scale tensors must be float16, bfloat16, or float32"
+        )
+        assert query_scale.dtype == key_scale.dtype == value_scale.dtype, (
+            "All scale tensors must have the same dtype"
+        )
+        assert query_scale.device == key_scale.device == value_scale.device == device, (
+            "All scale tensors must be on the same device as query/key/value"
+        )
+    else:
+        assert query.dtype in [torch.float16, torch.bfloat16], (
+            "Input dtype must be float16 or bfloat16"
+        )
+    assert query.dtype == key.dtype == value.dtype, (
+        "All inputs must have the same dtype"
+    )
+    assert num_heads_q % num_heads_kv == 0, (
+        "num_heads_q must be divisible by num_heads_kv"
+    )
+    assert head_dim % 16 == 0, (
+        "head_dim must be a multiple of 16 for efficient memory access"
+    )
+    assert head_dim <= 512, (
+        "head_dim must be less than or equal to 512 for efficient memory access"
+    )
+    if alpha is not None and delta is not None:
+        assert device == alpha.device == delta.device, (
+            "All inputs must be on the same device"
+        )
+    if cu_seqlens_q is not None and cu_seqlens_k is not None:
+        assert device == cu_seqlens_q.device == cu_seqlens_k.device, (
+            "All inputs must be on the same device"
+        )
+        assert cu_seqlens_q.dtype == cu_seqlens_k.dtype == torch.int32, (
+            "cu_seqlens_q and cu_seqlens_k must be int32"
+        )
+    if seqused_q is not None and seqused_k is not None:
+        assert device == seqused_q.device == seqused_k.device, (
+            "All inputs must be on the same device"
+        )
+        assert seqused_q.dtype == seqused_k.dtype == torch.int32, (
+            "seqused_q and seqused_k must be int32"
+        )
+    if window_sizes is not None:
+        assert window_sizes.dtype == torch.int32, "window_sizes must be int32"
+        assert window_sizes.ndim == 2 and window_sizes.shape[1] == 4, (
+            "window_sizes must have shape [num_kv_heads, 4] with columns [window_sink, window_left, window_right, window_dist]"
+        )
+def assert_bwd_inputs(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    dout: torch.Tensor,
+    lse: torch.Tensor,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    alpha: Optional[torch.Tensor] = None,
+    delta: Optional[torch.Tensor] = None,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    num_heads_q: int = None,
+    num_heads_kv: int = None,
+    head_dim: int = None,
+    is_quant: bool = False,
+    device: torch.device = None,
+):
+    """
+    Assert the validity of inputs for the backward base kernel.
+    :param query: Query tensor
+    :param key: Key tensor
+    :param value: Value tensor
+    :param out: Output tensor
+    :param dout: Gradient of the output tensor
+    :param lse: Log-sum-exp tensor
+    :param query_scale: Optional query scale tensor for quantized inputs
+    :param key_scale: Optional key scale tensor for quantized inputs
+    :param value_scale: Optional value scale tensor for quantized inputs
+    :param window_sizes: Optional window sizes tensor for local attention
+    :param alpha: Alpha tensor for gated attention
+    :param delta: Delta tensor for gated attention
+    :param cu_seqlens_q: Cumulative sequence lengths for queries
+    :param cu_seqlens_k: Cumulative sequence lengths for keys
+    :param seqused_q: Sequence used for queries
+    :param seqused_k: Sequence used for keys
+    :param num_heads_q: Number of query heads
+    :param num_heads_kv: Number of key/value heads
+    :param head_dim: Head dimension
+    :param is_quant: Whether the inputs are quantized
+    :param device: Device of the tensors
+    :raises AssertionError: If any of the assertions fail
+    """
+    assert (
+        device
+        == query.device
+        == key.device
+        == value.device
+        == out.device
+        == dout.device
+        == lse.device
+    ), "All inputs must be on the same device"
+    if is_quant:
+        assert query.dtype in [torch.float8_e5m2, torch.float8_e4m3fn], (
+            "Input dtype must be float8_e5m2 or float8_e4m3fn"
+        )
+        assert (
+            query_scale is not None
+            and key_scale is not None
+            and value_scale is not None
+        ), "Scale tensors must be provided for quantized inputs"
+        assert query_scale.dtype in [torch.float16, torch.bfloat16, torch.float32], (
+            "Scale tensors must be float16, bfloat16, or float32"
+        )
+        assert query_scale.dtype == key_scale.dtype == value_scale.dtype, (
+            "All scale tensors must have the same dtype"
+        )
+        assert query_scale.device == key_scale.device == value_scale.device == device, (
+            "All scale tensors must be on the same device as query/key/value"
+        )
+        assert query.dtype == key.dtype == value.dtype, (
+            "All inputs must have the same dtype"
+        )
+    else:
+        assert query.dtype in [torch.float16, torch.bfloat16], (
+            "Input dtype must be float16 or bfloat16"
+        )
+        assert query.dtype == key.dtype == value.dtype == out.dtype == dout.dtype, (
+            "All inputs must have the same dtype"
+        )
+    assert lse.dtype == torch.float32, (
+        "lse must be float32 for numerical stability in backward pass"
+    )
+    assert num_heads_q % num_heads_kv == 0, (
+        "num_heads_q must be divisible by num_heads_kv"
+    )
+    assert head_dim % 16 == 0, (
+        "head_dim must be a multiple of 16 for efficient memory access"
+    )
+    assert head_dim <= 512, (
+        "head_dim must be less than or equal to 512 for efficient memory access"
+    )
+    if alpha is not None and delta is not None:
+        assert device == alpha.device == delta.device, (
+            "All inputs must be on the same device"
+        )
+    if cu_seqlens_q is not None and cu_seqlens_k is not None:
+        assert device == cu_seqlens_q.device == cu_seqlens_k.device, (
+            "All inputs must be on the same device"
+        )
+        assert cu_seqlens_q.dtype == cu_seqlens_k.dtype == torch.int32, (
+            "cu_seqlens_q and cu_seqlens_k must be int32"
+        )
+    if seqused_q is not None and seqused_k is not None:
+        assert device == seqused_q.device == seqused_k.device, (
+            "All inputs must be on the same device"
+        )
+        assert seqused_q.dtype == seqused_k.dtype == torch.int32, (
+            "seqused_q and seqused_k must be int32"
+        )
+    if window_sizes is not None:
+        assert window_sizes.dtype == torch.int32, "window_sizes must be int32"
+        assert window_sizes.ndim == 2 and window_sizes.shape[1] == 4, (
+            "window_sizes must have shape [num_kv_heads, 4] with columns [window_sink, window_left, window_right, window_dist]"
+        )
+def assert_dec_inputs(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    alpha: Optional[torch.Tensor] = None,
+    delta: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    num_heads_q: int = None,
+    num_heads_kv: int = None,
+    head_dim: int = None,
+    is_quant: bool = False,
+    device: torch.device = None,
+):
+    """
+    Assert the validity of inputs for the decode kernel.
+    :param query: Query tensor
+    :param key: Key tensor
+    :param value: Value tensor
+    :param query_scale: Optional query scale tensor for quantized inputs
+    :param key_scale: Optional key scale tensor for quantized inputs
+    :param value_scale: Optional value scale tensor for quantized inputs
+    :param window_sizes: Optional window sizes tensor for local attention
+    :param alpha: Alpha tensor for gated attention
+    :param delta: Delta tensor for gated attention
+    :param cu_seqlens_k: Cumulative sequence lengths for keys
+    :param seqused_k: Sequence used for keys
+    :param num_heads_q: Number of query heads
+    :param num_heads_kv: Number of key/value heads
+    :param head_dim: Head dimension
+    :param is_quant: Whether the inputs are quantized
+    :param device: Device of the tensors
+    :raises AssertionError: If any of the assertions fail
+    """
+    assert device == query.device == key.device == value.device, (
+        "All inputs must be on the same device"
+    )
+    if is_quant:
+        assert query.dtype in [torch.float8_e5m2, torch.float8_e4m3fn], (
+            "Input dtype must be float8_e5m2 or float8_e4m3fn"
+        )
+        assert (
+            query_scale is not None
+            and key_scale is not None
+            and value_scale is not None
+        ), "Scale tensors must be provided for quantized inputs"
+        assert query_scale.dtype in [torch.float16, torch.bfloat16, torch.float32], (
+            "Scale tensors must be float16, bfloat16, or float32"
+        )
+        assert query_scale.dtype == key_scale.dtype == value_scale.dtype, (
+            "All scale tensors must have the same dtype"
+        )
+        assert query_scale.device == key_scale.device == value_scale.device == device, (
+            "All scale tensors must be on the same device as query/key/value"
+        )
+    else:
+        assert query.dtype in [torch.float16, torch.bfloat16], (
+            "Input dtype must be float16 or bfloat16"
+        )
+    assert query.dtype == key.dtype == value.dtype, (
+        "All inputs must have the same dtype"
+    )
+    assert num_heads_q % num_heads_kv == 0, (
+        "num_heads_q must be divisible by num_heads_kv"
+    )
+    assert head_dim % 16 == 0, (
+        "head_dim must be a multiple of 16 for efficient memory access"
+    )
+    assert head_dim <= 512, (
+        "head_dim must be less than or equal to 512 for efficient memory access"
+    )
+    if alpha is not None and delta is not None:
+        assert device == alpha.device == delta.device, (
+            "All inputs must be on the same device"
+        )
+    if cu_seqlens_k is not None:
+        assert device == cu_seqlens_k.device, "All inputs must be on the same device"
+        assert cu_seqlens_k.dtype == torch.int32, "cu_seqlens_k must be int32"
+    if seqused_k is not None:
+        assert device == seqused_k.device, "All inputs must be on the same device"
+        assert seqused_k.dtype == torch.int32, "seqused_k must be int32"
+    if window_sizes is not None:
+        assert window_sizes.dtype == torch.int32, "window_sizes must be int32"
+        assert window_sizes.ndim == 2 and window_sizes.shape[1] == 4, (
+            "window_sizes must have shape [num_kv_heads, 4] with columns [window_sink, window_left, window_right, window_dist]"
+        )

build/torch-cuda/autotuner.py ADDED Viewed

	@@ -0,0 +1,360 @@

+import torch
+import triton
+def _get_max_shared_mem():
+    props = torch.cuda.get_device_properties(torch.cuda.current_device())
+    return getattr(
+        props, "shared_memory_per_block_optin", props.shared_memory_per_block
+    )
+def _smem_bytes_fwd(tile_m, tile_n, tile_k, num_stages, dtype_bytes):
+    resident = tile_m * tile_k * dtype_bytes
+    pipelined = 2 * tile_n * tile_k * dtype_bytes * num_stages
+    return resident + pipelined
+def _smem_bytes_bwd(tile_m, tile_n, tile_k, num_stages, dtype_bytes):
+    resident = 2 * tile_n * tile_k * dtype_bytes
+    pipelined = 2 * tile_m * tile_k * dtype_bytes * num_stages
+    return resident + pipelined
+def _smem_bytes_dec(tile_m, tile_n, tile_k, num_stages, dtype_bytes):
+    resident = tile_m * tile_k * dtype_bytes
+    pipelined = 2 * tile_n * tile_k * dtype_bytes * num_stages
+    return resident + pipelined
+def _prune_fwd_configs(configs, named_args, **kwargs):
+    tile_k = kwargs.get("TILE_K", named_args.get("TILE_K", 128))
+    dtype_bytes = named_args["Q"].element_size()
+    max_smem = _get_max_shared_mem() - 4 * 1024
+    pruned = []
+    for cfg in configs:
+        tm, tn, ns = cfg.kwargs["TILE_M"], cfg.kwargs["TILE_N"], cfg.num_stages
+        if _smem_bytes_fwd(tm, tn, tile_k, ns, dtype_bytes) <= max_smem:
+            pruned.append(cfg)
+    if not pruned:
+        pruned = [
+            min(
+                configs,
+                key=lambda c: _smem_bytes_fwd(
+                    c.kwargs["TILE_M"],
+                    c.kwargs["TILE_N"],
+                    tile_k,
+                    c.num_stages,
+                    dtype_bytes,
+                ),
+            )
+        ]
+    return pruned
+def _prune_bwd_configs(configs, named_args, **kwargs):
+    tile_k = kwargs.get("TILE_K", named_args.get("TILE_K", 128))
+    dtype_bytes = named_args["Q"].element_size()
+    max_smem = _get_max_shared_mem() - 4 * 1024
+    pruned = []
+    for cfg in configs:
+        tm, tn, ns = cfg.kwargs["TILE_M"], cfg.kwargs["TILE_N"], cfg.num_stages
+        if _smem_bytes_bwd(tm, tn, tile_k, ns, dtype_bytes) <= max_smem:
+            pruned.append(cfg)
+    if not pruned:
+        pruned = [
+            min(
+                configs,
+                key=lambda c: _smem_bytes_bwd(
+                    c.kwargs["TILE_M"],
+                    c.kwargs["TILE_N"],
+                    tile_k,
+                    c.num_stages,
+                    dtype_bytes,
+                ),
+            )
+        ]
+    return pruned
+def _prune_dec_configs(configs, named_args, **kwargs):
+    tile_k = kwargs.get("TILE_K", named_args.get("TILE_K", 128))
+    seqlen_q = kwargs.get("seqlen_q", named_args.get("seqlen_q", None))
+    dtype_bytes = named_args["Q"].element_size()
+    max_smem = _get_max_shared_mem() - 4 * 1024
+    pruned = []
+    for cfg in configs:
+        tm, tn, ns = cfg.kwargs["TILE_M"], cfg.kwargs["TILE_N"], cfg.num_stages
+        if seqlen_q is not None and tm > seqlen_q:
+            continue
+        if _smem_bytes_dec(tm, tn, tile_k, ns, dtype_bytes) <= max_smem:
+            pruned.append(cfg)
+    if not pruned:
+        pruned = [
+            min(
+                configs,
+                key=lambda c: _smem_bytes_dec(
+                    c.kwargs["TILE_M"],
+                    c.kwargs["TILE_N"],
+                    tile_k,
+                    c.num_stages,
+                    dtype_bytes,
+                ),
+            )
+        ]
+    return pruned
+def get_fwd_dense_autotune_configs():
+    configs = []
+    for tile_m in [64, 128, 256]:
+        for tile_n in [32, 64, 128]:
+            for num_warps in [4, 8]:
+                for num_stages in [1, 2, 3]:
+                    configs.append(
+                        triton.Config(
+                            {"TILE_M": tile_m, "TILE_N": tile_n},
+                            num_warps=num_warps,
+                            num_stages=num_stages,
+                            num_ctas=1,
+                        )
+                    )
+    return configs
+def get_fwd_sparse_autotune_configs():
+    configs = []
+    for tile_m in [64, 128, 256]:
+        for tile_n in [32, 64, 128]:
+            for num_warps in [4, 8]:
+                for num_stages in [1, 2, 3]:
+                    configs.append(
+                        triton.Config(
+                            {"TILE_M": tile_m, "TILE_N": tile_n},
+                            num_warps=num_warps,
+                            num_stages=num_stages,
+                            num_ctas=1,
+                        )
+                    )
+    return configs
+def get_fwd_gated_autotune_configs():
+    configs = []
+    for tile_m in [64, 128, 256]:
+        for tile_n in [32, 64, 128]:
+            for num_warps in [4, 8]:
+                for num_stages in [1, 2, 3]:
+                    configs.append(
+                        triton.Config(
+                            {"TILE_M": tile_m, "TILE_N": tile_n},
+                            num_warps=num_warps,
+                            num_stages=num_stages,
+                            num_ctas=1,
+                        )
+                    )
+    return configs
+def get_bwd_dense_autotune_configs():
+    configs = []
+    for tile_m in [32, 64, 128]:
+        for tile_n in [64, 128, 256]:
+            for num_warps in [4, 8]:
+                for num_stages in [1, 2, 3]:
+                    configs.append(
+                        triton.Config(
+                            {"TILE_M": tile_m, "TILE_N": tile_n},
+                            num_warps=num_warps,
+                            num_stages=num_stages,
+                            num_ctas=1,
+                        )
+                    )
+    return configs
+def get_bwd_sparse_autotune_configs():
+    configs = []
+    for tile_m in [32, 64, 128]:
+        for tile_n in [64, 128, 256]:
+            for num_warps in [4, 8]:
+                for num_stages in [1, 2, 3]:
+                    configs.append(
+                        triton.Config(
+                            {"TILE_M": tile_m, "TILE_N": tile_n},
+                            num_warps=num_warps,
+                            num_stages=num_stages,
+                            num_ctas=1,
+                        )
+                    )
+    return configs
+def get_bwd_gated_autotune_configs():
+    configs = []
+    for tile_m in [32, 64, 128]:
+        for tile_n in [64, 128, 256]:
+            for num_warps in [4, 8]:
+                for num_stages in [1, 2, 3]:
+                    configs.append(
+                        triton.Config(
+                            {"TILE_M": tile_m, "TILE_N": tile_n},
+                            num_warps=num_warps,
+                            num_stages=num_stages,
+                            num_ctas=1,
+                        )
+                    )
+    return configs
+def get_dec_dense_autotune_configs():
+    configs = []
+    for tile_m in [16, 32, 64]:
+        for tile_n in [64, 128, 256]:
+            for num_warps in [4, 8]:
+                for num_stages in [1, 2, 3]:
+                    configs.append(
+                        triton.Config(
+                            {"TILE_M": tile_m, "TILE_N": tile_n},
+                            num_warps=num_warps,
+                            num_stages=num_stages,
+                            num_ctas=1,
+                        )
+                    )
+    return configs
+def get_dec_sparse_autotune_configs():
+    configs = []
+    for tile_m in [16, 32, 64]:
+        for tile_n in [64, 128, 256]:
+            for num_warps in [4, 8]:
+                for num_stages in [1, 2, 3]:
+                    configs.append(
+                        triton.Config(
+                            {"TILE_M": tile_m, "TILE_N": tile_n},
+                            num_warps=num_warps,
+                            num_stages=num_stages,
+                            num_ctas=1,
+                        )
+                    )
+    return configs
+def get_dec_gated_autotune_configs():
+    configs = []
+    for tile_m in [16, 32, 64]:
+        for tile_n in [64, 128, 256]:
+            for num_warps in [4, 8]:
+                for num_stages in [1, 2, 3]:
+                    configs.append(
+                        triton.Config(
+                            {"TILE_M": tile_m, "TILE_N": tile_n},
+                            num_warps=num_warps,
+                            num_stages=num_stages,
+                            num_ctas=1,
+                        )
+                    )
+    return configs
+def make_fwd_dense_autotuned_kernel(jit_kernel):
+    configs = get_fwd_dense_autotune_configs()
+    return triton.autotune(
+        configs=configs,
+        key=["SEQLEN_Q_CACHE", "SEQLEN_K_CACHE", "TILE_K", "IS_CAUSAL", "IS_LOCAL"],
+        prune_configs_by={"early_config_prune": _prune_fwd_configs},
+    )(jit_kernel)
+def make_fwd_sparse_autotuned_kernel(jit_kernel):
+    configs = get_fwd_sparse_autotune_configs()
+    return triton.autotune(
+        configs=configs,
+        key=["SEQLEN_Q_CACHE", "SEQLEN_K_CACHE", "TILE_K", "IS_CAUSAL", "IS_LOCAL"],
+        prune_configs_by={"early_config_prune": _prune_fwd_configs},
+    )(jit_kernel)
+def make_fwd_gated_autotuned_kernel(jit_kernel):
+    configs = get_fwd_gated_autotune_configs()
+    return triton.autotune(
+        configs=configs,
+        key=["SEQLEN_Q_CACHE", "SEQLEN_K_CACHE", "TILE_K", "IS_CAUSAL", "IS_LOCAL"],
+        prune_configs_by={"early_config_prune": _prune_fwd_configs},
+    )(jit_kernel)
+def make_bwd_dense_autotuned_kernel(jit_kernel):
+    configs = get_bwd_dense_autotune_configs()
+    return triton.autotune(
+        configs=configs,
+        key=["SEQLEN_Q_CACHE", "SEQLEN_K_CACHE", "TILE_K", "IS_CAUSAL", "IS_LOCAL"],
+        prune_configs_by={"early_config_prune": _prune_bwd_configs},
+    )(jit_kernel)
+def make_bwd_sparse_autotuned_kernel(jit_kernel):
+    configs = get_bwd_sparse_autotune_configs()
+    return triton.autotune(
+        configs=configs,
+        key=["SEQLEN_Q_CACHE", "SEQLEN_K_CACHE", "TILE_K", "IS_CAUSAL", "IS_LOCAL"],
+        prune_configs_by={"early_config_prune": _prune_bwd_configs},
+    )(jit_kernel)
+def make_bwd_gated_autotuned_kernel(jit_kernel):
+    configs = get_bwd_gated_autotune_configs()
+    return triton.autotune(
+        configs=configs,
+        key=["SEQLEN_Q_CACHE", "SEQLEN_K_CACHE", "TILE_K", "IS_CAUSAL", "IS_LOCAL"],
+        prune_configs_by={"early_config_prune": _prune_bwd_configs},
+    )(jit_kernel)
+def make_dec_dense_autotuned_kernel(jit_kernel):
+    configs = get_dec_dense_autotune_configs()
+    return triton.autotune(
+        configs=configs,
+        key=["SEQLEN_Q_CACHE", "SEQLEN_K_CACHE", "TILE_K", "IS_LOCAL"],
+        prune_configs_by={"early_config_prune": _prune_dec_configs},
+    )(jit_kernel)
+def make_dec_sparse_autotuned_kernel(jit_kernel):
+    configs = get_dec_sparse_autotune_configs()
+    return triton.autotune(
+        configs=configs,
+        key=["SEQLEN_Q_CACHE", "SEQLEN_K_CACHE", "TILE_K", "IS_LOCAL"],
+        prune_configs_by={"early_config_prune": _prune_dec_configs},
+    )(jit_kernel)
+def make_dec_gated_autotuned_kernel(jit_kernel):
+    configs = get_dec_gated_autotune_configs()
+    return triton.autotune(
+        configs=configs,
+        key=["SEQLEN_Q_CACHE", "SEQLEN_K_CACHE", "TILE_K", "IS_LOCAL"],
+        prune_configs_by={"early_config_prune": _prune_dec_configs},
+    )(jit_kernel)
+class AutotunedKernel:
+    STRIP_KWARGS = {"TILE_M", "TILE_N", "num_warps", "num_stages", "num_ctas"}
+    def __init__(self, autotuned_kernel):
+        self._autotuned = autotuned_kernel
+    def __getitem__(self, grid):
+        autotuned = self._autotuned
+        class _Launcher:
+            def __call__(_, *args, **kwargs):
+                for key in AutotunedKernel.STRIP_KWARGS:
+                    kwargs.pop(key, None)
+                return autotuned[grid](*args, **kwargs)
+        return _Launcher()
+    def __getattr__(self, name):
+        return getattr(self._autotuned, name)

build/torch-cuda/block_info.py ADDED Viewed

	@@ -0,0 +1,277 @@

+import triton
+import triton.language as tl
+@triton.jit
+def get_n_block_min_max(
+    seqlen_q,
+    seqlen_k,
+    m_block,
+    split_idx,
+    num_splits,
+    window_size_sink,
+    window_size_left,
+    window_size_right,
+    window_size_dist,
+    TILE_N: tl.constexpr,
+    TILE_M: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    IS_LOCAL: tl.constexpr,
+    IS_SPLIT_KV: tl.constexpr,
+    QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr,
+):
+    n_block_max = tl.cdiv(seqlen_k, TILE_N)
+    n_block_min = 0
+    n_block_window_max = n_block_max
+    n_block_window_min = n_block_min
+    n_block_sink_min = 0
+    n_block_sink_max = 0
+    if IS_CAUSAL or IS_LOCAL:
+        m_idx_max = (m_block + 1) * TILE_M
+        if QHEAD_PER_KVHEAD_PACKGQA > 1:
+            m_idx_max = tl.cdiv(m_idx_max, QHEAD_PER_KVHEAD_PACKGQA)
+        m_idx_max = tl.minimum(m_idx_max, seqlen_q)
+        n_idx = m_idx_max + seqlen_k - seqlen_q
+        n_block_max = tl.minimum(n_block_max, tl.cdiv(n_idx, TILE_N))
+        if IS_LOCAL:
+            n_idx_right = n_idx - window_size_dist - window_size_right
+            n_block_window_max = tl.minimum(
+                n_block_window_max, tl.maximum(tl.cdiv(n_idx_right, TILE_N), 0)
+            )
+    if IS_LOCAL:
+        n_block_sink_max = tl.minimum(
+            tl.cdiv(window_size_sink, TILE_N),
+            tl.maximum(tl.cdiv(n_idx, TILE_N), 0),
+        )
+        n_block_sink_exclude_max = tl.cdiv(window_size_sink, TILE_N)
+        m_idx_min = m_block * TILE_M
+        if QHEAD_PER_KVHEAD_PACKGQA > 1:
+            m_idx_min = m_idx_min // QHEAD_PER_KVHEAD_PACKGQA
+        n_idx = m_idx_min + seqlen_k - seqlen_q
+        n_idx_dist = n_idx - window_size_dist
+        n_block_min = tl.maximum(n_idx_dist // TILE_N, 0)
+        n_block_min = tl.maximum(n_block_min, n_block_sink_exclude_max)
+        n_idx_left = n_idx - window_size_dist - window_size_right - window_size_left
+        n_block_window_min = tl.maximum(n_idx_left // TILE_N, 0)
+        n_block_window_min = tl.maximum(n_block_window_min, n_block_sink_exclude_max)
+    if IS_SPLIT_KV:
+        if IS_LOCAL:
+            n_block_min = n_block_window_min
+            n_block_max_with_diag = n_block_max
+            n_block_max = tl.maximum(n_block_window_max, n_block_min)
+        total_n_blocks = tl.maximum(n_block_max - n_block_min, 0)
+        base = total_n_blocks // num_splits
+        extra = total_n_blocks % num_splits
+        n_block_min_new = n_block_min + tl.where(
+            split_idx < extra,
+            split_idx * (base + 1),
+            extra * (base + 1) + (split_idx - extra) * base,
+        )
+        n_block_count = tl.where(split_idx < extra, base + 1, base)
+        n_block_max = tl.minimum(n_block_min_new + n_block_count, n_block_max)
+        n_block_min = n_block_min_new
+        if IS_LOCAL:
+            n_block_max = tl.where(
+                split_idx >= num_splits - 1,
+                n_block_max_with_diag,
+                n_block_max,
+            )
+        if IS_SPLIT_KV:
+            n_block_sink_max = tl.where(split_idx == 0, n_block_sink_max, 0)
+    return (
+        n_block_min,
+        n_block_max,
+        n_block_window_min,
+        n_block_window_max,
+        n_block_sink_min,
+        n_block_sink_max,
+    )
+@triton.jit
+def get_m_block_min_max(
+    seqlen_q,
+    seqlen_k,
+    n_block,
+    split_idx,
+    num_splits,
+    window_size_sink,
+    window_size_left,
+    window_size_right,
+    window_size_dist,
+    TILE_N: tl.constexpr,
+    TILE_M: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    IS_LOCAL: tl.constexpr,
+    IS_SPLIT_QO: tl.constexpr,
+):
+    m_block_max = tl.cdiv(seqlen_q, TILE_M)
+    m_block_min = 0
+    m_block_window_max = m_block_max
+    m_block_window_min = m_block_min
+    m_block_sink_min = 0
+    m_block_sink_max = 0
+    if IS_CAUSAL or IS_LOCAL:
+        n_idx_min = n_block * TILE_N
+        m_idx = n_idx_min + seqlen_q - seqlen_k
+        m_block_min = tl.maximum(m_block_min, m_idx // TILE_M)
+        if IS_LOCAL:
+            m_idx_right = m_idx + window_size_dist + window_size_right
+            m_block_window_min = tl.maximum(m_block_window_min, m_idx_right // TILE_M)
+    if IS_LOCAL:
+        n_block_sink_exclude_max = tl.cdiv(window_size_sink, TILE_N)
+        is_sink_block = n_block < n_block_sink_exclude_max
+        n_idx_min = n_block * TILE_N
+        m_idx_sink = n_idx_min + seqlen_q - seqlen_k
+        m_block_sink_min = tl.maximum(m_idx_sink // TILE_M, 0)
+        m_block_sink_max = tl.where(is_sink_block, tl.cdiv(seqlen_q, TILE_M), 0)
+        m_block_sink_min = tl.where(is_sink_block, m_block_sink_min, 0)
+        n_idx_max = (n_block + 1) * TILE_N
+        m_idx = n_idx_max + seqlen_q - seqlen_k
+        m_idx_dist = m_idx + window_size_dist
+        m_block_max = tl.minimum(m_block_max, tl.cdiv(m_idx_dist, TILE_M))
+        m_idx_left = m_idx + window_size_dist + window_size_right + window_size_left
+        m_block_window_max = tl.minimum(m_block_window_max, tl.cdiv(m_idx_left, TILE_M))
+        m_block_min = tl.where(is_sink_block, 0, m_block_min)
+        m_block_max = tl.where(is_sink_block, 0, m_block_max)
+        m_block_window_min = tl.where(is_sink_block, 0, m_block_window_min)
+        m_block_window_max = tl.where(is_sink_block, 0, m_block_window_max)
+    if IS_SPLIT_QO:
+        if IS_LOCAL:
+            n_idx_max = (n_block + 1) * TILE_N
+            m_idx = n_idx_max + seqlen_q - seqlen_k
+            m_block_min_no_mask = tl.maximum(m_block_min, tl.cdiv(m_idx, TILE_M))
+            m_block_window_min = tl.maximum(m_block_window_min, m_block_min_no_mask)
+            win_blocks = tl.maximum(m_block_window_max - m_block_window_min, 0)
+            base = win_blocks // num_splits
+            extra = win_blocks % num_splits
+            win_off = tl.where(
+                split_idx < extra,
+                split_idx * (base + 1),
+                extra * (base + 1) + (split_idx - extra) * base,
+            )
+            win_cnt = tl.where(split_idx < extra, base + 1, base)
+            m_block_window_min = m_block_window_min + win_off
+            m_block_window_max = tl.minimum(
+                m_block_window_min + win_cnt, m_block_window_max
+            )
+            # First split keeps diagonal, others skip it
+            m_block_min = tl.where(split_idx == 0, m_block_min, m_block_max)
+            m_block_sink_max = tl.where(split_idx == 0, m_block_sink_max, 0)
+            m_block_sink_min = tl.where(split_idx == 0, m_block_sink_min, 0)
+        else:
+            total_m_blocks = tl.maximum(m_block_max - m_block_min, 0)
+            base = total_m_blocks // num_splits
+            extra = total_m_blocks % num_splits
+            m_block_min_new = m_block_min + tl.where(
+                split_idx < extra,
+                split_idx * (base + 1),
+                extra * (base + 1) + (split_idx - extra) * base,
+            )
+            m_block_count = tl.where(split_idx < extra, base + 1, base)
+            m_block_max = tl.minimum(m_block_min_new + m_block_count, m_block_max)
+            m_block_min = m_block_min_new
+    return (
+        m_block_min,
+        m_block_max,
+        m_block_window_min,
+        m_block_window_max,
+        m_block_sink_min,
+        m_block_sink_max,
+    )
+@triton.jit
+def get_n_block_min_causal_local_mask(
+    seqlen_q,
+    seqlen_k,
+    m_block,
+    n_block_min,
+    window_size_right,
+    window_size_dist,
+    TILE_N: tl.constexpr,
+    TILE_M: tl.constexpr,
+    IS_LOCAL: tl.constexpr,
+    QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr,
+):
+    m_idx_min = m_block * TILE_M
+    if QHEAD_PER_KVHEAD_PACKGQA > 1:
+        m_idx_min = m_idx_min // QHEAD_PER_KVHEAD_PACKGQA
+    n_idx = m_idx_min + seqlen_k - seqlen_q
+    n_idx_right = (
+        n_idx if not IS_LOCAL else n_idx - window_size_dist - window_size_right
+    )
+    return tl.maximum(n_block_min, n_idx_right // TILE_N)
+@triton.jit
+def get_n_block_min_before_local_mask(
+    seqlen_q,
+    seqlen_k,
+    m_block,
+    n_block_min,
+    window_size_left,
+    window_size_right,
+    window_size_dist,
+    TILE_N: tl.constexpr,
+    TILE_M: tl.constexpr,
+    IS_LOCAL: tl.constexpr,
+    QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr,
+):
+    if not IS_LOCAL:
+        return n_block_min
+    else:
+        m_idx_max = (m_block + 1) * TILE_M
+        if QHEAD_PER_KVHEAD_PACKGQA > 1:
+            m_idx_max = tl.cdiv(m_idx_max, QHEAD_PER_KVHEAD_PACKGQA)
+        m_idx_max = tl.minimum(m_idx_max, seqlen_q)
+        n_idx = m_idx_max + seqlen_k - seqlen_q
+        n_idx_left = n_idx - window_size_dist - window_size_right - window_size_left
+        return tl.maximum(n_block_min, tl.cdiv(n_idx_left, TILE_N))
+@triton.jit
+def get_m_block_min_causal_local_mask(
+    seqlen_q,
+    seqlen_k,
+    n_block,
+    m_block_min,
+    window_size_right,
+    window_size_dist,
+    TILE_N: tl.constexpr,
+    TILE_M: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    IS_LOCAL: tl.constexpr,
+):
+    if not IS_CAUSAL and not IS_LOCAL:
+        return m_block_min
+    else:
+        n_idx_max = (n_block + 1) * TILE_N
+        m_idx = n_idx_max + seqlen_q - seqlen_k
+        if IS_LOCAL:
+            m_idx_right = m_idx + window_size_dist + window_size_right
+        else:
+            m_idx_right = m_idx
+        return tl.maximum(m_block_min, tl.cdiv(m_idx_right, TILE_M))
+@triton.jit
+def get_m_block_max_before_local_mask(
+    seqlen_q,
+    seqlen_k,
+    n_block,
+    m_block_max,
+    window_size_left,
+    window_size_right,
+    window_size_dist,
+    TILE_N: tl.constexpr,
+    TILE_M: tl.constexpr,
+    IS_LOCAL: tl.constexpr,
+):
+    if not IS_LOCAL:
+        return m_block_max
+    else:
+        n_idx_min = n_block * TILE_N
+        m_idx = n_idx_min + seqlen_q - seqlen_k
+        m_idx_left = m_idx + window_size_dist + window_size_right + window_size_left
+        return tl.minimum(m_block_max, m_idx_left // TILE_M)

build/torch-cuda/cache_utils.py ADDED Viewed

	@@ -0,0 +1,180 @@

+import functools
+import inspect
+from collections import OrderedDict
+import torch
+@functools.lru_cache(maxsize=8)
+def get_device_num_sms(device: torch.device) -> int:
+    """
+    Get the SM count for a given device.
+    :param device: Torch device.
+    :return num_sms: Number of streaming multiprocessors.
+    """
+    return torch.cuda.get_device_properties(device).multi_processor_count
+@functools.lru_cache(maxsize=8)
+def get_device_arch(device: torch.device) -> int:
+    """
+    Get the architecture for a given device.
+    :param device: Torch device.
+    :return arch: Architecture model as a number.
+    """
+    if device.type == "cuda":
+        major, minor = torch.cuda.get_device_capability(device)
+        sm = major * 10 + minor
+        return sm if sm >= 80 else -1
+    if device.type in {"xpu", "mps", "cpu"}:
+        return -1
+    raise ValueError(f"Unsupported device: {device}")
+def cache_launch_config(fn):
+    """
+    Cache launch-config lookups by (arch, *args, **kwargs).
+    :param fn: Launch-config function to wrap.
+    :return wrapper: Wrapped function with cache.
+    """
+    params = list(inspect.signature(fn).parameters.keys())
+    device_idx = params.index("device") if "device" in params else None
+    arch_idx = params.index("arch") if "arch" in params else None
+    cached = functools.lru_cache(maxsize=None)(
+        lambda _arch, *args, **kwargs: fn(*args, **kwargs)
+    )
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        device = (
+            args[device_idx]
+            if (device_idx is not None and device_idx < len(args))
+            else kwargs.get("device")
+        )
+        arch = (
+            args[arch_idx]
+            if (arch_idx is not None and arch_idx < len(args))
+            else kwargs.get("arch")
+        )
+        return cached((device.type, arch), *args, **kwargs)
+    return wrapper
+def cache_launch_grid(fn, maxsize: int = 512):
+    """
+    Bounded cache for Triton grid-factory closures.
+    :param fn: Grid-factory function to wrap.
+    :param maxsize: Maximum number of cached entries.
+    :return wrapper: Wrapped function.
+    """
+    return functools.lru_cache(maxsize=maxsize)(fn)
+class CachedKernel:
+    """
+    Wrap a `triton.jit`-decorated function so that each launch reuses a
+    previously-compiled `CompiledKernel` when the specialization matches.
+    """
+    __slots__ = ("_kernel", "_cache")
+    def __init__(self, kernel):
+        self._kernel = kernel
+        self._cache = OrderedDict()
+    def __getitem__(self, grid):
+        kernel = self._kernel
+        kernel_warmup = kernel.warmup
+        kernel_arg_names = kernel.arg_names
+        cache = self._cache
+        maxsize = 1024
+        tensor_indices = None
+        static_grid = None
+        if not callable(grid):
+            static_grid = tuple(grid)
+            if len(static_grid) < 3:
+                static_grid = static_grid + (1,) * (3 - len(static_grid))
+        def launcher(
+            *args,
+            num_warps: int = 4,
+            num_stages: int = 1,
+            num_ctas: int = 1,
+            **constexprs,
+        ):
+            nonlocal tensor_indices
+            if static_grid is not None:
+                launch_grid = static_grid
+            else:
+                launch_grid = tuple(grid(constexprs))
+                if len(launch_grid) < 3:
+                    launch_grid = launch_grid + (1,) * (3 - len(launch_grid))
+            if tensor_indices is None:
+                tensor_indices = tuple(
+                    i for i, a in enumerate(args) if torch.is_tensor(a)
+                )
+            key = (
+                tuple((args[i].dtype, args[i].device.type) for i in tensor_indices),
+                tuple(constexprs.values()) if constexprs else (),
+                num_warps,
+                num_stages,
+                num_ctas,
+            )
+            compiled = cache.get(key)
+            if compiled is None:
+                compiled = kernel_warmup(
+                    *args,
+                    grid=launch_grid,
+                    num_warps=num_warps,
+                    num_stages=num_stages,
+                    num_ctas=num_ctas,
+                    **constexprs,
+                )
+                cache[key] = compiled
+                if len(cache) > maxsize:
+                    cache.popitem(last=False)
+            else:
+                cache.move_to_end(key)
+            runner = compiled[launch_grid]
+            if not constexprs:
+                runner(*args)
+                return
+            args_iter = iter(args)
+            full_args = [
+                constexprs[name] if name in constexprs else next(args_iter)
+                for name in kernel_arg_names
+            ]
+            runner(*full_args)
+        return launcher
+    def __getattr__(self, name):
+        return getattr(self._kernel, name)
+def wrap_kernel(kernel):
+    """
+    Wrap a `triton.jit` kernel for compiled-kernel caching.
+    :param kernel: Triton JIT kernel.
+    :return wrapped: Wrapped kernel preserving the `kernel[grid](...)` syntax.
+    """
+    return CachedKernel(kernel)

build/torch-cuda/flash_bwd_postprocess.py ADDED Viewed

	@@ -0,0 +1,215 @@

+from typing import Optional
+import torch
+import triton
+import triton.language as tl
+from . import (
+    utils,
+    cache_utils,
+    launch_grid,
+    seqlen_info,
+    kernel_repr,
+)
+@triton.jit(repr=kernel_repr.bwd_postprocess_repr)
+def _bwd_postprocess_kernel(
+    dQaccum,
+    dQ,
+    dAaccum,
+    dA,
+    stride_dqab,
+    stride_dqah,
+    stride_dqam,
+    stride_dqb,
+    stride_dqh,
+    stride_dqm,
+    stride_daab,
+    stride_daah,
+    stride_daam,
+    stride_dab,
+    stride_dah,
+    stride_dam,
+    cu_seqlens_q,
+    seqused_q,
+    seqlen_q,
+    head_dim,
+    head_dim_rounded,
+    scale,
+    HAS_CU_SEQLENS_Q: tl.constexpr,
+    HAS_SEQUSED_Q: tl.constexpr,
+    HAS_DA: tl.constexpr,
+    TILE_M: tl.constexpr,
+    TILE_K: tl.constexpr,
+):
+    m_block = tl.program_id(0)
+    head_idx = tl.program_id(1)
+    batch_idx = tl.program_id(2)
+    # Get seqlen info for this batch
+    offset_q, actual_seqlen_q = seqlen_info.get_seqlen_info(
+        batch_idx=batch_idx,
+        seqlen_static=seqlen_q,
+        cu_seqlens=cu_seqlens_q,
+        seqused=seqused_q,
+        HAS_CU_SEQLENS=HAS_CU_SEQLENS_Q,
+        HAS_SEQUSED=HAS_SEQUSED_Q,
+    )
+    padded_offset_q = (offset_q + batch_idx * TILE_M) // TILE_M * TILE_M
+    # Initialize base pointers
+    dq_accum_base = seqlen_info.offset_batch_Q(
+        dQaccum + head_idx * stride_dqah,
+        batch_idx,
+        offset_q,
+        padded_offset_q,
+        stride_dqab,
+        stride_dqam,
+        HAS_CU_SEQLENS_Q,
+        USE_PADDED=True,
+    )
+    dq_base = seqlen_info.offset_batch_Q(
+        dQ + head_idx * stride_dqh,
+        batch_idx,
+        offset_q,
+        padded_offset_q,
+        stride_dqb,
+        stride_dqm,
+        HAS_CU_SEQLENS_Q,
+        USE_PADDED=False,
+    )
+    # Create descriptors
+    dq_accum_desc = tl.make_tensor_descriptor(
+        base=dq_accum_base,
+        shape=[actual_seqlen_q, head_dim_rounded],
+        strides=[stride_dqam, 1],
+        block_shape=[TILE_M, TILE_K],
+    )
+    dq_desc = tl.make_tensor_descriptor(
+        base=dq_base,
+        shape=[actual_seqlen_q, head_dim],
+        strides=[stride_dqm, 1],
+        block_shape=[TILE_M, TILE_K],
+    )
+    # Load accumulators
+    acc_dq = dq_accum_desc.load([m_block * TILE_M, 0])
+    # Scale dq
+    dq = (acc_dq * scale).to(dQ.dtype.element_ty)
+    # Store dq
+    dq_desc.store([m_block * TILE_M, 0], dq)
+    if HAS_DA:
+        da_accum_base = seqlen_info.offset_batch_Q(
+            dAaccum + head_idx * stride_daah,
+            batch_idx,
+            offset_q,
+            padded_offset_q,
+            stride_daab,
+            stride_daam,
+            HAS_CU_SEQLENS_Q,
+            USE_PADDED=True,
+        )
+        da_base = seqlen_info.offset_batch_Q(
+            dA + head_idx * stride_dah,
+            batch_idx,
+            offset_q,
+            padded_offset_q,
+            stride_dab,
+            stride_dam,
+            HAS_CU_SEQLENS_Q,
+            USE_PADDED=False,
+        )
+        da_accum_desc = tl.make_tensor_descriptor(
+            base=da_accum_base,
+            shape=[actual_seqlen_q],
+            strides=[stride_daam],
+            block_shape=[TILE_M],
+        )
+        da_desc = tl.make_tensor_descriptor(
+            base=da_base,
+            shape=[actual_seqlen_q],
+            strides=[stride_dam],
+            block_shape=[TILE_M],
+        )
+        # Load da accumulators
+        acc_da = da_accum_desc.load([m_block * TILE_M])
+        # Scale da
+        da = (acc_da * scale).to(dA.dtype.element_ty)
+        # Store da
+        da_desc.store([m_block * TILE_M], da)
+_bwd_postprocess_kernel = cache_utils.wrap_kernel(_bwd_postprocess_kernel)
+def _flash_attn_bwd_postprocess(
+    dq_accum: torch.Tensor,
+    dq: torch.Tensor,
+    scale: float,
+    head_dim_rounded: int,
+    da_accum: Optional[torch.Tensor] = None,
+    da: Optional[torch.Tensor] = None,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    tile_m: int = 128,
+    tile_k: int = 128,
+) -> torch.Tensor:
+    is_varlen = cu_seqlens_q is not None
+    if not is_varlen:
+        batch_size, seqlen_q, num_heads_q, head_dim = dq.shape
+    else:
+        _, num_heads_q, head_dim = dq.shape
+        batch_size = cu_seqlens_q.shape[0] - 1
+        seqlen_q = max_seqlen_q
+    has_da = da_accum is not None and da is not None
+    grid = launch_grid.get_bwd_postprocess_grid(batch_size, seqlen_q, num_heads_q)
+    triton.set_allocator(utils.alloc_fn)
+    _bwd_postprocess_kernel[grid](
+        dq_accum,
+        dq,
+        da_accum if has_da else dq_accum,
+        da if has_da else dq,
+        dq_accum.stride(0) if not is_varlen else 0,
+        dq_accum.stride(1) if not is_varlen else dq_accum.stride(0),
+        head_dim_rounded,
+        dq.stride(0) if not is_varlen else 0,
+        dq.stride(-2),
+        dq.stride(-3) if not is_varlen else dq.stride(0),
+        da_accum.stride(0) if (has_da and not is_varlen) else 0,
+        da_accum.stride(1)
+        if (has_da and not is_varlen)
+        else (da_accum.stride(0) if has_da else 0),
+        da_accum.stride(-1) if has_da else 0,
+        da.stride(0) if (has_da and not is_varlen) else 0,
+        da.stride(-2) if has_da else 0,
+        da.stride(-1) if has_da else 0,
+        cu_seqlens_q,
+        seqused_q,
+        seqlen_q,
+        head_dim,
+        head_dim_rounded,
+        scale,
+        HAS_CU_SEQLENS_Q=is_varlen,
+        HAS_SEQUSED_Q=seqused_q is not None,
+        HAS_DA=has_da,
+        TILE_M=tile_m,
+        TILE_K=tile_k,
+        num_warps=4,
+        num_stages=1,
+        num_ctas=1,
+    )
+    return dq

build/torch-cuda/flash_bwd_preprocess.py ADDED Viewed

	@@ -0,0 +1,266 @@

+from typing import Optional
+import math
+import torch
+import triton
+import triton.language as tl
+from . import (
+    utils,
+    cache_utils,
+    launch_grid,
+    seqlen_info,
+    activations,
+    kernel_repr,
+)
+@triton.jit(repr=kernel_repr.bwd_preprocess_repr)
+def _bwd_preprocess_kernel(
+    Out,
+    dO,
+    dPsum,
+    LSE,
+    LSELog2,
+    dQaccum,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_pb,
+    stride_ph,
+    stride_pm,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_l2b,
+    stride_l2h,
+    stride_l2m,
+    stride_dqab,
+    stride_dqah,
+    stride_dqam,
+    cu_seqlens_q,
+    seqused_q,
+    seqlen_q,
+    head_dim,
+    head_dim_rounded,
+    HAS_CU_SEQLENS_Q: tl.constexpr,
+    HAS_SEQUSED_Q: tl.constexpr,
+    TILE_M: tl.constexpr,
+    TILE_K: tl.constexpr,
+):
+    m_block = tl.program_id(0)
+    head_idx = tl.program_id(1)
+    batch_idx = tl.program_id(2)
+    # Get seqlen info for this batch
+    offset_q, actual_seqlen_q = seqlen_info.get_seqlen_info(
+        batch_idx=batch_idx,
+        seqlen_static=seqlen_q,
+        cu_seqlens=cu_seqlens_q,
+        seqused=seqused_q,
+        HAS_CU_SEQLENS=HAS_CU_SEQLENS_Q,
+        HAS_SEQUSED=HAS_SEQUSED_Q,
+    )
+    padded_offset_q = (offset_q + batch_idx * TILE_M) // TILE_M * TILE_M
+    # Initialize base pointers
+    o_base = seqlen_info.offset_batch_Q(
+        Out + head_idx * stride_oh,
+        batch_idx,
+        offset_q,
+        padded_offset_q,
+        stride_ob,
+        stride_om,
+        HAS_CU_SEQLENS_Q,
+        USE_PADDED=False,
+    )
+    do_base = seqlen_info.offset_batch_Q(
+        dO + head_idx * stride_doh,
+        batch_idx,
+        offset_q,
+        padded_offset_q,
+        stride_dob,
+        stride_dom,
+        HAS_CU_SEQLENS_Q,
+        USE_PADDED=False,
+    )
+    dpsum_base = seqlen_info.offset_batch_Q(
+        dPsum + head_idx * stride_ph,
+        batch_idx,
+        offset_q,
+        padded_offset_q,
+        stride_pb,
+        stride_pm,
+        HAS_CU_SEQLENS_Q,
+        USE_PADDED=True,
+    )
+    lse_base = seqlen_info.offset_batch_Q(
+        LSE + head_idx * stride_lh,
+        batch_idx,
+        offset_q,
+        padded_offset_q,
+        stride_lb,
+        stride_lm,
+        HAS_CU_SEQLENS_Q,
+        USE_PADDED=False,
+    )
+    lse_log2_base = seqlen_info.offset_batch_Q(
+        LSELog2 + head_idx * stride_l2h,
+        batch_idx,
+        offset_q,
+        padded_offset_q,
+        stride_l2b,
+        stride_l2m,
+        HAS_CU_SEQLENS_Q,
+        USE_PADDED=True,
+    )
+    dq_accum_base = seqlen_info.offset_batch_Q(
+        dQaccum + head_idx * stride_dqah,
+        batch_idx,
+        offset_q,
+        padded_offset_q,
+        stride_dqab,
+        stride_dqam,
+        HAS_CU_SEQLENS_Q,
+        USE_PADDED=True,
+    )
+    # Create descriptors
+    o_desc = tl.make_tensor_descriptor(
+        base=o_base,
+        shape=[actual_seqlen_q, head_dim],
+        strides=[stride_om, 1],
+        block_shape=[TILE_M, TILE_K],
+    )
+    do_desc = tl.make_tensor_descriptor(
+        base=do_base,
+        shape=[actual_seqlen_q, head_dim],
+        strides=[stride_dom, 1],
+        block_shape=[TILE_M, TILE_K],
+    )
+    dpsum_desc = tl.make_tensor_descriptor(
+        base=dpsum_base,
+        shape=[actual_seqlen_q],
+        strides=[stride_pm],
+        block_shape=[TILE_M],
+    )
+    lse_desc = tl.make_tensor_descriptor(
+        base=lse_base,
+        shape=[actual_seqlen_q],
+        strides=[stride_lm],
+        block_shape=[TILE_M],
+    )
+    lse_log2_desc = tl.make_tensor_descriptor(
+        base=lse_log2_base,
+        shape=[actual_seqlen_q],
+        strides=[stride_l2m],
+        block_shape=[TILE_M],
+    )
+    dq_accum_desc = tl.make_tensor_descriptor(
+        base=dq_accum_base,
+        shape=[actual_seqlen_q, head_dim_rounded],
+        strides=[stride_dqam, 1],
+        block_shape=[TILE_M, TILE_K],
+    )
+    # Initialize accumulators
+    acc_dq = tl.zeros((TILE_M, TILE_K), dtype=tl.float32)
+    # Load o tile
+    o_tile = o_desc.load([m_block * TILE_M, 0]).to(tl.float32)
+    # Load do tile
+    do_tile = do_desc.load([m_block * TILE_M, 0]).to(tl.float32)
+    # Compute dpsum
+    dpsum = tl.sum(o_tile * do_tile, axis=1)
+    # Store dpsum
+    dpsum_desc.store([m_block * TILE_M], dpsum)
+    # Load lse tile
+    lse_tile = lse_desc.load([m_block * TILE_M])
+    # Compute lse_log2
+    lse_log2 = tl.where(lse_tile == lse_tile, lse_tile * math.log2(math.e), 1e6)
+    lse_log2 = activations.check_inf(lse_log2)
+    # Store dq_accum
+    dq_accum_desc.store([m_block * TILE_M, 0], acc_dq)
+    # Store lse_log2
+    lse_log2_desc.store([m_block * TILE_M], lse_log2)
+_bwd_preprocess_kernel = cache_utils.wrap_kernel(_bwd_preprocess_kernel)
+def _flash_attn_bwd_preprocess(
+    out: torch.Tensor,
+    dout: torch.Tensor,
+    dpsum: torch.Tensor,
+    lse: torch.Tensor,
+    lse_log2: torch.Tensor,
+    dq_accum: torch.Tensor,
+    head_dim_rounded: int,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    tile_m: int = 128,
+    tile_k: int = 128,
+) -> torch.Tensor:
+    is_varlen = cu_seqlens_q is not None
+    if not is_varlen:
+        batch_size, seqlen_q, num_heads_q, head_dim = out.shape
+    else:
+        _, num_heads_q, head_dim = out.shape
+        batch_size = cu_seqlens_q.shape[0] - 1
+        seqlen_q = max_seqlen_q
+    grid = launch_grid.get_bwd_preprocess_grid(batch_size, seqlen_q, num_heads_q)
+    triton.set_allocator(utils.alloc_fn)
+    _bwd_preprocess_kernel[grid](
+        out,
+        dout,
+        dpsum,
+        lse,
+        lse_log2,
+        dq_accum,
+        out.stride(0) if not is_varlen else 0,
+        out.stride(-2),
+        out.stride(-3) if not is_varlen else out.stride(0),
+        dout.stride(0) if not is_varlen else 0,
+        dout.stride(-2),
+        dout.stride(-3) if not is_varlen else dout.stride(0),
+        dpsum.stride(0) if not is_varlen else 0,
+        dpsum.stride(1) if not is_varlen else dpsum.stride(0),
+        dpsum.stride(-1),
+        lse.stride(0) if not is_varlen else 0,
+        lse.stride(1) if not is_varlen else lse.stride(0),
+        lse.stride(-1),
+        lse_log2.stride(0) if not is_varlen else 0,
+        lse_log2.stride(1) if not is_varlen else lse_log2.stride(0),
+        lse_log2.stride(-1),
+        dq_accum.stride(0) if not is_varlen else 0,
+        dq_accum.stride(1) if not is_varlen else dq_accum.stride(0),
+        head_dim_rounded,
+        cu_seqlens_q,
+        seqused_q,
+        seqlen_q,
+        head_dim,
+        head_dim_rounded,
+        HAS_CU_SEQLENS_Q=is_varlen,
+        HAS_SEQUSED_Q=seqused_q is not None,
+        TILE_M=tile_m,
+        TILE_K=tile_k,
+        num_warps=4,
+        num_stages=1,
+        num_ctas=1,
+    )
+    return dpsum

build/torch-cuda/flash_dec_combine.py ADDED Viewed

	@@ -0,0 +1,222 @@

+import torch
+import triton
+import triton.language as tl
+from . import (
+    utils,
+    cache_utils,
+    launch_template,
+    launch_grid,
+    seqlen_info,
+    kernel_repr,
+)
+@triton.jit(repr=kernel_repr.dec_combine_repr)
+def _dec_combine_kernel(
+    Out_partial,
+    Lse_partial,
+    Out,
+    Lse,
+    stride_ops,
+    stride_opb,
+    stride_oph,
+    stride_opm,
+    stride_lps,
+    stride_lpb,
+    stride_lph,
+    stride_lpm,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    cu_seqlens_q,
+    seqused_q,
+    num_splits,
+    num_heads_q,
+    head_dim,
+    TILE_K: tl.constexpr,
+    HAS_CU_SEQLENS_Q: tl.constexpr,
+    HAS_SEQUSED_Q: tl.constexpr,
+):
+    bh_idx = tl.program_id(0)
+    batch_idx = bh_idx // num_heads_q
+    head_idx = bh_idx - batch_idx * num_heads_q
+    # Get seqlen info for this batch
+    offset_q, _ = seqlen_info.get_seqlen_info(
+        batch_idx=batch_idx,
+        seqlen_static=1,
+        cu_seqlens=cu_seqlens_q,
+        seqused=seqused_q,
+        HAS_CU_SEQLENS=HAS_CU_SEQLENS_Q,
+        HAS_SEQUSED=HAS_SEQUSED_Q,
+    )
+    # Initialize base pointers
+    out_part_base = seqlen_info.offset_batch_Q(
+        Out_partial + head_idx * stride_oph,
+        batch_idx,
+        offset_q,
+        0,
+        stride_opb,
+        stride_opm,
+        HAS_CU_SEQLENS_Q,
+        USE_PADDED=False,
+    )
+    lse_part_base = seqlen_info.offset_batch_Q(
+        Lse_partial + head_idx * stride_lph,
+        batch_idx,
+        offset_q,
+        0,
+        stride_lpb,
+        stride_lpm,
+        HAS_CU_SEQLENS_Q,
+        USE_PADDED=False,
+    )
+    out_base = seqlen_info.offset_batch_Q(
+        Out + head_idx * stride_oh,
+        batch_idx,
+        offset_q,
+        0,
+        stride_ob,
+        stride_om,
+        HAS_CU_SEQLENS_Q,
+        USE_PADDED=False,
+    )
+    lse_base = seqlen_info.offset_batch_Q(
+        Lse + head_idx * stride_lh,
+        batch_idx,
+        offset_q,
+        0,
+        stride_lb,
+        stride_lm,
+        HAS_CU_SEQLENS_Q,
+        USE_PADDED=False,
+    )
+    # Create descriptors
+    out_part_desc = tl.make_tensor_descriptor(
+        base=out_part_base,
+        shape=[num_splits, head_dim],
+        strides=[stride_ops, 1],
+        block_shape=[1, TILE_K],
+    )
+    out_desc = tl.make_tensor_descriptor(
+        base=out_base,
+        shape=[head_dim],
+        strides=[1],
+        block_shape=[TILE_K],
+    )
+    # Initialize accumulators
+    e_sum = tl.full((), 0.0, dtype=tl.float32)
+    e_max = tl.full((), float("-inf"), dtype=tl.float32)
+    acc_o = tl.zeros((TILE_K,), dtype=tl.float32)
+    # Combine split outputs
+    for s in tl.range(0, num_splits):
+        # Load partial LSE
+        lse_s = tl.load(lse_part_base + s * stride_lps).to(tl.float32)
+        # Compute normalized exponentials
+        new_e_max = tl.maximum(lse_s, e_max)
+        old_scale = tl.where(e_max == float("-inf"), 0.0, tl.exp2(e_max - new_e_max))
+        exp_logic = tl.where(lse_s == float("-inf"), 0.0, tl.exp2(lse_s - new_e_max))
+        # Load partial outputs
+        o_s = tl.sum(out_part_desc.load([s, 0]), axis=0)
+        # Compute scaled outputs
+        acc_o *= old_scale
+        acc_o += exp_logic * o_s
+        # Update e_sum and e_max
+        e_sum = e_sum * old_scale + exp_logic
+        e_max = new_e_max
+    # Normalize output
+    inv_sum = tl.where((e_sum == 0.0) | (e_sum != e_sum), 0.0, 1.0 / e_sum)
+    acc_o *= inv_sum
+    # Store output
+    out_desc.store([0], acc_o.to(Out.dtype.element_ty))
+    # Compute LSE
+    # ln2 = math.log(2.0)
+    ln2 = 0.6931471805599453
+    lse = tl.where(e_sum > 0.0, (e_max + tl.log2(e_sum)) * ln2, float("-inf"))
+    # Store LSE
+    tl.store(lse_base, lse)
+_dec_combine_kernel = cache_utils.wrap_kernel(_dec_combine_kernel)
+def _flash_attn_dec_combine(
+    out_partial: torch.Tensor,
+    lse_partial: torch.Tensor,
+    out: torch.Tensor,
+    lse: torch.Tensor,
+    cu_seqlens_q: torch.Tensor = None,
+    seqused_q: torch.Tensor = None,
+):
+    device = out.device
+    arch = cache_utils.get_device_arch(device)
+    is_varlen = cu_seqlens_q is not None
+    num_splits = out_partial.shape[0]
+    if not is_varlen:
+        batch_size, num_heads_q, head_dim = out_partial.shape[1:]
+    else:
+        _, num_heads_q, head_dim = out_partial.shape[1:]
+        batch_size = cu_seqlens_q.shape[0] - 1
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    num_warps, num_stages, num_ctas = launch_template.get_dec_combine_launch_config(
+        tile_k=TILE_K,
+        device=device,
+        arch=arch,
+    )
+    grid = launch_grid.get_dec_combine_grid(
+        batch_size=batch_size,
+        num_heads_q=num_heads_q,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    _dec_combine_kernel[grid](
+        out_partial,
+        lse_partial,
+        out,
+        lse,
+        out_partial.stride(0),
+        out_partial.stride(1) if not is_varlen else 0,
+        out_partial.stride(-2),
+        1 if not is_varlen else out_partial.stride(1),
+        lse_partial.stride(0),
+        lse_partial.stride(1) if not is_varlen else 0,
+        lse_partial.stride(-1),
+        1 if not is_varlen else lse_partial.stride(1),
+        out.stride(0) if not is_varlen else 0,
+        out.stride(-2),
+        1 if not is_varlen else out.stride(0),
+        lse.stride(0) if not is_varlen else 0,
+        lse.stride(-1),
+        1 if not is_varlen else lse.stride(0),
+        cu_seqlens_q,
+        seqused_q,
+        num_splits,
+        num_heads_q,
+        head_dim,
+        TILE_K=TILE_K,
+        HAS_CU_SEQLENS_Q=is_varlen,
+        HAS_SEQUSED_Q=seqused_q is not None,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )

build/torch-cuda/flash_dense_bwd.py ADDED Viewed

	@@ -0,0 +1,1019 @@

+from typing import Optional, Tuple
+import math
+import torch
+import triton
+import triton.language as tl
+from . import (
+    assert_inputs,
+    utils,
+    cache_utils,
+    launch_template,
+    launch_grid,
+    activations,
+    flash_bwd_preprocess,
+    flash_bwd_postprocess,
+    kernel_repr,
+    autotuner,
+)
+from .scheduler import (
+    AttnBwdGridIndex,
+    AttnBwdConfig,
+    AttnBwdBlockScheduler,
+    AttnBwdPointerScheduler,
+    AttnMaskScheduler,
+)
+@triton.jit
+def _bwd_inner_dense_kernel(
+    config,
+    ptrs_sched,
+    mask_sched,
+    acc_dk,
+    acc_dv,
+    k_tile,
+    v_tile,
+    q_ptrs,
+    do_ptrs,
+    dq_accum_ptrs,
+    lse_ptrs,
+    dpsum_ptrs,
+    m_block,
+    IS_MASK: tl.constexpr,
+    MASK_CAUSAL: tl.constexpr,
+    MASK_LOCAL: tl.constexpr,
+    MASK_SINK: tl.constexpr,
+):
+    # Load query tile
+    q_tile = ptrs_sched.load_q(config, q_ptrs, m_block)
+    # Compute attention scores
+    acc_s = tl.dot(k_tile, q_tile.T)
+    # Load LSE
+    lse_log2 = ptrs_sched.load_lse(config, lse_ptrs, m_block)
+    if IS_MASK:
+        # Apply mask to attention scores
+        acc_s = mask_sched.apply_mask(
+            acc_s=acc_s,
+            iter_block=m_block,
+            MASK_CAUSAL=MASK_CAUSAL,
+            MASK_LOCAL=MASK_LOCAL,
+            MASK_SINK=MASK_SINK,
+        )
+    # Compute attention weights
+    p = activations.exp2(acc_s * config.softmax_scale_log2 - lse_log2[None, :]).to(
+        q_tile.dtype
+    )
+    # Load output gradients tile
+    do_tile = ptrs_sched.load_do(config, do_ptrs, m_block)
+    # Compute value gradients
+    acc_dv += tl.dot(p, do_tile)
+    # Compute attention weight gradients
+    acc_dp = tl.dot(v_tile, tl.trans(do_tile))
+    # Load dpsum
+    dpsum = ptrs_sched.load_dpsum(config, dpsum_ptrs, m_block)
+    # Compute attention score gradients
+    ds = p * (acc_dp - dpsum[None, :]).to(q_tile.dtype)
+    # Compute query gradients
+    dq = tl.dot(tl.trans(ds), k_tile)
+    # Store query gradients
+    ptrs_sched.store_dq(config, dq_accum_ptrs, m_block, dq)
+    # Compute key gradients
+    acc_dk += tl.dot(ds, q_tile)
+    return acc_dk, acc_dv
+@triton.jit(repr=kernel_repr.bwd_dense_repr)
+def _bwd_dense_kernel(
+    Q,
+    K,
+    V,
+    dO,
+    LSELog2,
+    dPsum,
+    dQaccum,
+    dK,
+    dV,
+    softmax_scale,
+    query_scale,
+    key_scale,
+    value_scale,
+    window_sizes,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_lb,
+    stride_lh,
+    stride_pb,
+    stride_ph,
+    stride_dqab,
+    stride_dqah,
+    stride_dqam,
+    stride_dkb,
+    stride_dkh,
+    stride_dkn,
+    stride_dks,
+    stride_dvb,
+    stride_dvh,
+    stride_dvn,
+    stride_dvs,
+    stride_wh,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    seqused_q,
+    seqused_k,
+    num_splits,
+    seqlen_q,
+    seqlen_k,
+    head_dim,
+    SEQLEN_Q_CACHE: tl.constexpr,
+    SEQLEN_K_CACHE: tl.constexpr,
+    QHEAD_PER_KVHEAD: tl.constexpr,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+    TILE_K: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    IS_LOCAL: tl.constexpr,
+    IS_SPLIT_QO: tl.constexpr,
+    HAS_CU_SEQLENS_Q: tl.constexpr,
+    HAS_CU_SEQLENS_K: tl.constexpr,
+    HAS_SEQUSED_Q: tl.constexpr,
+    HAS_SEQUSED_K: tl.constexpr,
+):
+    # Create grid index
+    grid_idx = AttnBwdGridIndex.create(
+        num_splits=num_splits,
+        QHEAD_PER_KVHEAD=QHEAD_PER_KVHEAD,
+        IS_SPLIT_QO=IS_SPLIT_QO,
+    )
+    # Load window sizes
+    (
+        window_size_sink,
+        window_size_left,
+        window_size_right,
+        window_size_dist,
+    ) = grid_idx.load_window_sizes(
+        window_sizes=window_sizes,
+        stride_wh=stride_wh,
+        IS_LOCAL=IS_LOCAL,
+    )
+    # Create config
+    config = AttnBwdConfig.create(
+        softmax_scale=softmax_scale,
+        query_scale=query_scale,
+        key_scale=key_scale,
+        value_scale=value_scale,
+        n_block=grid_idx.n_block,
+        batch_idx=grid_idx.batch_idx,
+        window_size_sink=window_size_sink,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        window_size_dist=window_size_dist,
+        head_dim=head_dim,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        QHEAD_PER_KVHEAD=QHEAD_PER_KVHEAD,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_CAUSAL=IS_CAUSAL,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+        HAS_SEQUSED_Q=HAS_SEQUSED_Q,
+        HAS_SEQUSED_K=HAS_SEQUSED_K,
+    )
+    # Create pointer scheduler
+    ptrs_sched = AttnBwdPointerScheduler.create(
+        config=config,
+        Q=Q,
+        K=K,
+        V=V,
+        dO=dO,
+        LSELog2=LSELog2,
+        dPsum=dPsum,
+        dQaccum=dQaccum,
+        dK=dK,
+        dV=dV,
+        batch_idx=grid_idx.batch_idx,
+        head_idx=grid_idx.head_idx,
+        head_kv_idx=grid_idx.head_kv_idx,
+        split_idx=grid_idx.split_idx,
+        stride_qb=stride_qb,
+        stride_qh=stride_qh,
+        stride_qm=stride_qm,
+        stride_kb=stride_kb,
+        stride_kh=stride_kh,
+        stride_kn=stride_kn,
+        stride_vb=stride_vb,
+        stride_vh=stride_vh,
+        stride_vn=stride_vn,
+        stride_dob=stride_dob,
+        stride_doh=stride_doh,
+        stride_dom=stride_dom,
+        stride_lb=stride_lb,
+        stride_lh=stride_lh,
+        stride_pb=stride_pb,
+        stride_ph=stride_ph,
+        stride_dqab=stride_dqab,
+        stride_dqah=stride_dqah,
+        stride_dqam=stride_dqam,
+        stride_dkb=stride_dkb,
+        stride_dkh=stride_dkh,
+        stride_dkn=stride_dkn,
+        stride_dks=stride_dks,
+        stride_dvb=stride_dvb,
+        stride_dvh=stride_dvh,
+        stride_dvn=stride_dvn,
+        stride_dvs=stride_dvs,
+        IS_SPLIT_QO=IS_SPLIT_QO,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+    )
+    # Create block scheduler
+    block_sched = AttnBwdBlockScheduler.create(
+        config=config,
+        split_idx=grid_idx.split_idx,
+        num_splits=num_splits,
+        IS_CAUSAL=IS_CAUSAL,
+        IS_LOCAL=IS_LOCAL,
+        IS_SPLIT_QO=IS_SPLIT_QO,
+    )
+    # Create mask scheduler
+    mask_sched = AttnMaskScheduler.create(config, SWAP_AB=True)
+    # Early exit if no n_blocks to process
+    if grid_idx.n_block * TILE_N >= config.actual_seqlen_k:
+        return
+    # Create pointers
+    q_ptrs = ptrs_sched.make_q_ptrs(config)
+    k_ptrs = ptrs_sched.make_k_ptrs(config)
+    v_ptrs = ptrs_sched.make_v_ptrs(config)
+    do_ptrs = ptrs_sched.make_do_ptrs(config)
+    lse_ptrs = ptrs_sched.make_lse_ptrs(config)
+    dpsum_ptrs = ptrs_sched.make_dpsum_ptrs(config)
+    dq_accum_ptrs = ptrs_sched.make_dq_accum_ptrs(config)
+    dk_ptrs = ptrs_sched.make_dk_ptrs(config)
+    dv_ptrs = ptrs_sched.make_dv_ptrs(config)
+    # Initialize accumulators
+    acc_dk = tl.zeros((TILE_N, TILE_K), dtype=tl.float32)
+    acc_dv = tl.zeros((TILE_N, TILE_K), dtype=tl.float32)
+    # Load key tile
+    k_tile = ptrs_sched.load_k(config, k_ptrs)
+    # Load value tile
+    v_tile = ptrs_sched.load_v(config, v_ptrs)
+    # Process m_blocks with causal masking
+    if IS_CAUSAL or IS_LOCAL:
+        for m_block in tl.range(
+            block_sched.m_block_min, block_sched.m_block_min_no_mask
+        ):
+            acc_dk, acc_dv = _bwd_inner_dense_kernel(
+                config=config,
+                ptrs_sched=ptrs_sched,
+                mask_sched=mask_sched,
+                acc_dk=acc_dk,
+                acc_dv=acc_dv,
+                k_tile=k_tile,
+                v_tile=v_tile,
+                q_ptrs=q_ptrs,
+                do_ptrs=do_ptrs,
+                dq_accum_ptrs=dq_accum_ptrs,
+                lse_ptrs=lse_ptrs,
+                dpsum_ptrs=dpsum_ptrs,
+                m_block=m_block,
+                IS_MASK=True,
+                MASK_CAUSAL=IS_CAUSAL,
+                MASK_LOCAL=True if IS_LOCAL else False,
+                MASK_SINK=False,
+            )
+    # Process m_blocks without masking
+    if not IS_LOCAL and block_sched.m_block_min_no_mask < block_sched.m_block_max:
+        for m_block in tl.range(
+            block_sched.m_block_min_no_mask, block_sched.m_block_max
+        ):
+            acc_dk, acc_dv = _bwd_inner_dense_kernel(
+                config=config,
+                ptrs_sched=ptrs_sched,
+                mask_sched=mask_sched,
+                acc_dk=acc_dk,
+                acc_dv=acc_dv,
+                k_tile=k_tile,
+                v_tile=v_tile,
+                q_ptrs=q_ptrs,
+                do_ptrs=do_ptrs,
+                dq_accum_ptrs=dq_accum_ptrs,
+                lse_ptrs=lse_ptrs,
+                dpsum_ptrs=dpsum_ptrs,
+                m_block=m_block,
+                IS_MASK=False,
+                MASK_CAUSAL=False,
+                MASK_LOCAL=False,
+                MASK_SINK=False,
+            )
+    if IS_LOCAL:
+        # Process m_blocks with local right masking
+        if block_sched.m_block_window_min < block_sched.m_block_window_min_no_mask:
+            for m_block in tl.range(
+                block_sched.m_block_window_min,
+                block_sched.m_block_window_min_no_mask,
+            ):
+                acc_dk, acc_dv = _bwd_inner_dense_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    acc_dk=acc_dk,
+                    acc_dv=acc_dv,
+                    k_tile=k_tile,
+                    v_tile=v_tile,
+                    q_ptrs=q_ptrs,
+                    do_ptrs=do_ptrs,
+                    dq_accum_ptrs=dq_accum_ptrs,
+                    lse_ptrs=lse_ptrs,
+                    dpsum_ptrs=dpsum_ptrs,
+                    m_block=m_block,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                )
+        # Process m_blocks without masking
+        if (
+            block_sched.m_block_window_min_no_mask
+            < block_sched.m_block_window_max_no_mask
+        ):
+            for m_block in tl.range(
+                block_sched.m_block_window_min_no_mask,
+                block_sched.m_block_window_max_no_mask,
+            ):
+                acc_dk, acc_dv = _bwd_inner_dense_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    acc_dk=acc_dk,
+                    acc_dv=acc_dv,
+                    k_tile=k_tile,
+                    v_tile=v_tile,
+                    q_ptrs=q_ptrs,
+                    do_ptrs=do_ptrs,
+                    dq_accum_ptrs=dq_accum_ptrs,
+                    lse_ptrs=lse_ptrs,
+                    dpsum_ptrs=dpsum_ptrs,
+                    m_block=m_block,
+                    IS_MASK=False,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=False,
+                    MASK_SINK=False,
+                )
+        # Process m_blocks with local left masking
+        if block_sched.m_block_window_max_no_mask < block_sched.m_block_window_max:
+            for m_block in tl.range(
+                block_sched.m_block_window_max_no_mask,
+                block_sched.m_block_window_max,
+            ):
+                acc_dk, acc_dv = _bwd_inner_dense_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    acc_dk=acc_dk,
+                    acc_dv=acc_dv,
+                    k_tile=k_tile,
+                    v_tile=v_tile,
+                    q_ptrs=q_ptrs,
+                    do_ptrs=do_ptrs,
+                    dq_accum_ptrs=dq_accum_ptrs,
+                    lse_ptrs=lse_ptrs,
+                    dpsum_ptrs=dpsum_ptrs,
+                    m_block=m_block,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                )
+        # Process m_blocks with local sink masking
+        if block_sched.m_block_sink_max > block_sched.m_block_sink_min:
+            for m_block in tl.range(
+                block_sched.m_block_sink_min,
+                block_sched.m_block_sink_max,
+            ):
+                acc_dk, acc_dv = _bwd_inner_dense_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    acc_dk=acc_dk,
+                    acc_dv=acc_dv,
+                    k_tile=k_tile,
+                    v_tile=v_tile,
+                    q_ptrs=q_ptrs,
+                    do_ptrs=do_ptrs,
+                    dq_accum_ptrs=dq_accum_ptrs,
+                    lse_ptrs=lse_ptrs,
+                    dpsum_ptrs=dpsum_ptrs,
+                    m_block=m_block,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=True,
+                )
+    # Store value gradients
+    ptrs_sched.store_dv(config, dv_ptrs, acc_dv)
+    # Scale key gradients
+    acc_dk = acc_dk * softmax_scale
+    # Store key gradients
+    ptrs_sched.store_dk(config, dk_ptrs, acc_dk)
+_bwd_dense_kernel = cache_utils.wrap_kernel(_bwd_dense_kernel)
+_bwd_dense_kernel_autotuned = None
+def _get_autotuned_kernel():
+    global _bwd_dense_kernel_autotuned
+    if _bwd_dense_kernel_autotuned is None:
+        jit_kernel = _bwd_dense_kernel._kernel
+        autotuned = autotuner.make_bwd_dense_autotuned_kernel(jit_kernel)
+        _bwd_dense_kernel_autotuned = autotuner.AutotunedKernel(autotuned)
+    return _bwd_dense_kernel_autotuned
+def _flash_dense_attn_backward(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    dout: torch.Tensor,
+    lse: torch.Tensor,
+    is_causal: bool = False,
+    softmax_scale: float = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    is_split_qo: bool = False,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    batch_size, seqlen_q, num_heads_q, head_dim = query.shape
+    _, seqlen_k, num_heads_kv, _ = key.shape
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_bwd_inputs(
+            query,
+            key,
+            value,
+            out,
+            dout,
+            lse,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            cu_seqlens_q=None,
+            cu_seqlens_k=None,
+            seqused_q=None,
+            seqused_k=None,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name="bwd_dense",
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_causal=is_causal,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _bwd_dense_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        # Placeholder for pre-launch computations
+        TILE_M = TILE_N = 64
+        num_warps = num_stages = num_ctas = None
+    num_splits = (
+        utils.num_splits_heuristic(
+            seqlen_q=seqlen_k,
+            seqlen_k=seqlen_q,
+            num_SMs=num_SMs,
+            TILE_M=TILE_N,
+            TILE_N=TILE_M,
+            max_split_blocks=utils.max_split_blocks_from_window_sizes(
+                window_sizes, TILE_M
+            )
+            if is_local
+            else None,
+        )
+        if is_split_qo
+        else 1
+    )
+    seqlen_q_rounded = int(math.ceil(seqlen_q / 128) * 128)
+    head_dim_rounded = int(math.ceil(head_dim / 32) * 32)
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    dq = torch.empty_like(query, dtype=query_scale.dtype)
+    dk = torch.empty_like(key, dtype=key_scale.dtype)
+    dv = torch.empty_like(value, dtype=value_scale.dtype)
+    lse_log2 = torch.empty(
+        (batch_size, num_heads_q, seqlen_q_rounded),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dpsum = torch.empty(
+        (batch_size, num_heads_q, seqlen_q_rounded),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dq_accum = torch.empty(
+        (batch_size, num_heads_q, seqlen_q_rounded * head_dim_rounded),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dk_accum = torch.zeros(
+        (num_splits, batch_size, seqlen_k, num_heads_kv, head_dim)
+        if is_split_qo and num_splits > 1
+        else (batch_size, seqlen_k, num_heads_kv, head_dim),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dv_accum = torch.zeros(
+        (num_splits, batch_size, seqlen_k, num_heads_kv, head_dim)
+        if is_split_qo and num_splits > 1
+        else (batch_size, seqlen_k, num_heads_kv, head_dim),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    flash_bwd_preprocess._flash_attn_bwd_preprocess(
+        out=out,
+        dout=dout,
+        dpsum=dpsum,
+        lse=lse,
+        lse_log2=lse_log2,
+        dq_accum=dq_accum,
+        head_dim_rounded=head_dim_rounded,
+        tile_m=TILE_M,
+        tile_k=TILE_K,
+    )
+    grid = launch_grid.get_bwd_grid(
+        seqlen_k=seqlen_k,
+        num_heads_q=num_heads_q,
+        batch_size=batch_size,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        dout,
+        lse_log2,
+        dpsum,
+        dq_accum,
+        dk_accum,
+        dv_accum,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        query.stride(0),
+        query.stride(-2),
+        query.stride(-3),
+        key.stride(0),
+        key.stride(-2),
+        key.stride(-3),
+        value.stride(0),
+        value.stride(-2),
+        value.stride(-3),
+        dout.stride(0),
+        dout.stride(-2),
+        dout.stride(-3),
+        lse_log2.stride(0),
+        lse_log2.stride(-2),
+        dpsum.stride(0),
+        dpsum.stride(-2),
+        dq_accum.stride(0),
+        dq_accum.stride(1),
+        head_dim_rounded,
+        dk_accum.stride(-4) if is_split_qo and num_splits > 1 else dk_accum.stride(0),
+        dk_accum.stride(-2),
+        dk_accum.stride(-3),
+        dk_accum.stride(0) if is_split_qo and num_splits > 1 else 0,
+        dv_accum.stride(-4) if is_split_qo and num_splits > 1 else dv_accum.stride(0),
+        dv_accum.stride(-2),
+        dv_accum.stride(-3),
+        dv_accum.stride(0) if is_split_qo and num_splits > 1 else 0,
+        window_sizes.stride(0),
+        None,
+        None,
+        None,
+        None,
+        num_splits,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=seqlen_q // 1024,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD=qhead_per_kvhead,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_CAUSAL=is_causal,
+        IS_LOCAL=is_local,
+        IS_SPLIT_QO=is_split_qo and num_splits > 1,
+        HAS_CU_SEQLENS_Q=False,
+        HAS_CU_SEQLENS_K=False,
+        HAS_SEQUSED_Q=False,
+        HAS_SEQUSED_K=False,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name="bwd_dense",
+                seqlen_q=seqlen_q,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_causal=is_causal,
+                is_quant=is_quant,
+            )
+    flash_bwd_postprocess._flash_attn_bwd_postprocess(
+        dq_accum=dq_accum,
+        dq=dq,
+        scale=softmax_scale,
+        head_dim_rounded=head_dim_rounded,
+        tile_m=TILE_M,
+        tile_k=TILE_K,
+    )
+    if is_split_qo and num_splits > 1:
+        dk_accum = dk_accum.sum(dim=0)
+        dv_accum = dv_accum.sum(dim=0)
+    dk.copy_(dk_accum)
+    dv.copy_(dv_accum)
+    return dq, dk, dv
+def _flash_dense_attn_varlen_backward(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    dout: torch.Tensor,
+    lse: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    is_causal: bool = False,
+    softmax_scale: float = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    is_split_qo: bool = False,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    total_q, num_heads_q, head_dim = query.shape
+    total_k, num_heads_kv, _ = key.shape
+    batch_size = cu_seqlens_q.shape[0] - 1
+    seqlen_q = max_seqlen_q
+    seqlen_k = max_seqlen_k
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_bwd_inputs(
+            query,
+            key,
+            value,
+            out,
+            dout,
+            lse,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name="bwd_dense",
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_causal=is_causal,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _bwd_dense_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        # Placeholder for pre-launch computations
+        TILE_M = TILE_N = 64
+        num_warps = num_stages = num_ctas = None
+    num_splits = (
+        utils.num_splits_heuristic(
+            seqlen_q=seqlen_k,
+            seqlen_k=seqlen_q,
+            num_SMs=num_SMs,
+            TILE_M=TILE_N,
+            TILE_N=TILE_M,
+            max_split_blocks=utils.max_split_blocks_from_window_sizes(
+                window_sizes, TILE_M
+            )
+            if is_local
+            else None,
+        )
+        if is_split_qo
+        else 1
+    )
+    total_q_rounded_padded = int(math.ceil((total_q + batch_size * 128) / 128) * 128)
+    head_dim_rounded = int(math.ceil(head_dim / 32) * 32)
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    dq = torch.empty_like(query, dtype=query_scale.dtype)
+    dk = torch.empty_like(key, dtype=key_scale.dtype)
+    dv = torch.empty_like(value, dtype=value_scale.dtype)
+    lse_log2 = torch.empty(
+        num_heads_q,
+        total_q_rounded_padded,
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dpsum = torch.empty(
+        num_heads_q,
+        total_q_rounded_padded,
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dq_accum = torch.empty(
+        num_heads_q,
+        total_q_rounded_padded * head_dim_rounded,
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dk_accum = torch.zeros(
+        (num_splits, total_k, num_heads_kv, head_dim)
+        if is_split_qo and num_splits > 1
+        else (total_k, num_heads_kv, head_dim),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dv_accum = torch.zeros(
+        (num_splits, total_k, num_heads_kv, head_dim)
+        if is_split_qo and num_splits > 1
+        else (total_k, num_heads_kv, head_dim),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    flash_bwd_preprocess._flash_attn_bwd_preprocess(
+        out=out,
+        dout=dout,
+        dpsum=dpsum,
+        lse=lse,
+        lse_log2=lse_log2,
+        dq_accum=dq_accum,
+        head_dim_rounded=head_dim_rounded,
+        cu_seqlens_q=cu_seqlens_q,
+        seqused_q=seqused_q,
+        max_seqlen_q=max_seqlen_q,
+        tile_m=TILE_M,
+        tile_k=TILE_K,
+    )
+    grid = launch_grid.get_bwd_grid(
+        seqlen_k=seqlen_k,
+        num_heads_q=num_heads_q,
+        batch_size=batch_size,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        dout,
+        lse_log2,
+        dpsum,
+        dq_accum,
+        dk_accum,
+        dv_accum,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        0,
+        query.stride(-2),
+        query.stride(0),
+        0,
+        key.stride(-2),
+        key.stride(0),
+        0,
+        value.stride(-2),
+        value.stride(0),
+        0,
+        dout.stride(-2),
+        dout.stride(0),
+        0,
+        lse_log2.stride(0),
+        0,
+        dpsum.stride(0),
+        0,
+        dq_accum.stride(0),
+        head_dim_rounded,
+        0,
+        dk_accum.stride(-2),
+        dk_accum.stride(-3),
+        dk_accum.stride(0) if is_split_qo and num_splits > 1 else 0,
+        0,
+        dv_accum.stride(-2),
+        dv_accum.stride(-3),
+        dv_accum.stride(0) if is_split_qo and num_splits > 1 else 0,
+        window_sizes.stride(0),
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_q,
+        seqused_k,
+        num_splits,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=seqlen_q // 1024,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD=qhead_per_kvhead,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_CAUSAL=is_causal,
+        IS_LOCAL=is_local,
+        IS_SPLIT_QO=is_split_qo and num_splits > 1,
+        HAS_CU_SEQLENS_Q=True,
+        HAS_CU_SEQLENS_K=True,
+        HAS_SEQUSED_Q=seqused_q is not None,
+        HAS_SEQUSED_K=seqused_k is not None,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name="bwd_dense",
+                seqlen_q=seqlen_q,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_causal=is_causal,
+                is_quant=is_quant,
+            )
+    flash_bwd_postprocess._flash_attn_bwd_postprocess(
+        dq_accum=dq_accum,
+        dq=dq,
+        scale=softmax_scale,
+        head_dim_rounded=head_dim_rounded,
+        cu_seqlens_q=cu_seqlens_q,
+        seqused_q=seqused_q,
+        max_seqlen_q=max_seqlen_q,
+        tile_m=TILE_M,
+        tile_k=TILE_K,
+    )
+    if is_split_qo and num_splits > 1:
+        dk_accum = dk_accum.sum(dim=0)
+        dv_accum = dv_accum.sum(dim=0)
+    dk.copy_(dk_accum)
+    dv.copy_(dv_accum)
+    return dq, dk, dv

build/torch-cuda/flash_dense_dec.py ADDED Viewed

	@@ -0,0 +1,857 @@

+from typing import Tuple, Optional
+import torch
+import triton
+import triton.language as tl
+from . import (
+    assert_inputs,
+    utils,
+    cache_utils,
+    launch_template,
+    launch_grid,
+    flash_dec_combine,
+    kernel_repr,
+    autotuner,
+)
+from .scheduler import (
+    AttnDecGridIndex,
+    AttnDecConfig,
+    AttnDecBlockScheduler,
+    AttnDecPointerScheduler,
+    AttnMaskScheduler,
+    SoftmaxScheduler,
+)
+@triton.jit
+def _dec_inner_dense_kernel(
+    config: AttnDecConfig,
+    ptrs_sched: AttnDecPointerScheduler,
+    mask_sched: AttnMaskScheduler,
+    softmax_sched: SoftmaxScheduler,
+    q_tile,
+    k_tile,
+    k_ptrs,
+    v_ptrs,
+    acc_o,
+    row_max,
+    row_sum,
+    n_block,
+    n_block_min,
+    IS_MASK: tl.constexpr,
+    MASK_LOCAL: tl.constexpr,
+    MASK_SINK: tl.constexpr,
+    CHECK_INF: tl.constexpr,
+):
+    # Compute attention scores
+    acc_s = tl.dot(q_tile, k_tile.T)
+    # Prefetch next key tile
+    if n_block > n_block_min:
+        # Load next key tile
+        k_tile = ptrs_sched.load_k(config, k_ptrs, n_block - 1)
+    if IS_MASK:
+        # Apply mask to attention scores
+        acc_s = mask_sched.apply_mask(
+            acc_s=acc_s,
+            iter_block=n_block,
+            MASK_LOCAL=MASK_LOCAL,
+            MASK_SINK=MASK_SINK,
+        )
+    # Apply online softmax
+    p, row_max, row_sum, row_scale = softmax_sched.online_softmax(
+        acc_s=acc_s,
+        row_max=row_max,
+        row_sum=row_sum,
+        CHECK_INF=CHECK_INF,
+    )
+    # Load value tile
+    v_tile = ptrs_sched.load_v(config, v_ptrs, n_block)
+    # Rescale output accumulator
+    acc_o = softmax_sched.rescale_o(
+        acc_o=acc_o,
+        row_scale=row_scale,
+    )
+    # Update output accumulator
+    acc_o += tl.dot(p.to(v_tile.dtype), v_tile)
+    return k_tile, acc_o, row_max, row_sum
+@triton.jit(repr=kernel_repr.dec_dense_repr)
+def _dec_dense_kernel(
+    Q,
+    K,
+    V,
+    Out,
+    Lse,
+    softmax_scale,
+    query_scale,
+    key_scale,
+    value_scale,
+    window_sizes,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_os,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_ls,
+    stride_wh,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    seqused_q,
+    seqused_k,
+    num_splits,
+    seqlen_q,
+    seqlen_k,
+    head_dim,
+    SEQLEN_Q_CACHE: tl.constexpr,
+    SEQLEN_K_CACHE: tl.constexpr,
+    QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+    TILE_K: tl.constexpr,
+    IS_LOCAL: tl.constexpr,
+    HAS_CU_SEQLENS_Q: tl.constexpr,
+    HAS_CU_SEQLENS_K: tl.constexpr,
+    HAS_SEQUSED_Q: tl.constexpr,
+    HAS_SEQUSED_K: tl.constexpr,
+):
+    # Create grid index
+    grid_idx = AttnDecGridIndex.create(
+        num_splits=num_splits,
+    )
+    # Load window sizes
+    (
+        window_size_sink,
+        window_size_left,
+        window_size_right,
+        window_size_dist,
+    ) = grid_idx.load_window_sizes(
+        window_sizes=window_sizes,
+        stride_wh=stride_wh,
+        IS_LOCAL=IS_LOCAL,
+    )
+    # Create config
+    config = AttnDecConfig.create(
+        softmax_scale=softmax_scale,
+        query_scale=query_scale,
+        key_scale=key_scale,
+        value_scale=value_scale,
+        batch_idx=grid_idx.batch_idx,
+        window_size_sink=window_size_sink,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        window_size_dist=window_size_dist,
+        head_dim=head_dim,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        QHEAD_PER_KVHEAD_PACKGQA=QHEAD_PER_KVHEAD_PACKGQA,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+        HAS_SEQUSED_Q=HAS_SEQUSED_Q,
+        HAS_SEQUSED_K=HAS_SEQUSED_K,
+    )
+    # Create pointer scheduler
+    ptrs_sched = AttnDecPointerScheduler.create(
+        config=config,
+        Q=Q,
+        K=K,
+        V=V,
+        Out=Out,
+        Lse=Lse,
+        batch_idx=grid_idx.batch_idx,
+        head_idx=grid_idx.head_idx,
+        head_kv_idx=grid_idx.head_kv_idx,
+        split_idx=grid_idx.split_idx,
+        stride_qb=stride_qb,
+        stride_qh=stride_qh,
+        stride_qm=stride_qm,
+        stride_kb=stride_kb,
+        stride_kh=stride_kh,
+        stride_kn=stride_kn,
+        stride_vb=stride_vb,
+        stride_vh=stride_vh,
+        stride_vn=stride_vn,
+        stride_ob=stride_ob,
+        stride_oh=stride_oh,
+        stride_om=stride_om,
+        stride_os=stride_os,
+        stride_lb=stride_lb,
+        stride_lh=stride_lh,
+        stride_lm=stride_lm,
+        stride_ls=stride_ls,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+    )
+    # Create block scheduler
+    block_sched = AttnDecBlockScheduler.create(
+        config=config,
+        split_idx=grid_idx.split_idx,
+        num_splits=num_splits,
+        IS_LOCAL=IS_LOCAL,
+    )
+    # Create mask scheduler
+    mask_sched = AttnMaskScheduler.create(config)
+    # Create softmax scheduler
+    softmax_sched = SoftmaxScheduler.create(config)
+    # Create pointers
+    out_ptrs = ptrs_sched.make_out_ptrs(config)
+    lse_ptrs = ptrs_sched.make_lse_ptrs(config)
+    # Early exit if no n_blocks to process
+    if block_sched.is_empty():
+        ptrs_sched.store_empty(config, out_ptrs, lse_ptrs, Out)
+        return
+    q_ptrs = ptrs_sched.make_q_ptrs(config)
+    k_ptrs = ptrs_sched.make_k_ptrs(config)
+    v_ptrs = ptrs_sched.make_v_ptrs(config)
+    # Initialize accumulators
+    row_max = tl.full((TILE_M,), float("-inf"), dtype=tl.float32)
+    row_sum = tl.zeros((TILE_M,), dtype=tl.float32)
+    acc_o = tl.zeros((TILE_M, TILE_K), dtype=tl.float32)
+    # Load query tile
+    q_tile = ptrs_sched.load_q(config, q_ptrs)
+    # Load key tile
+    k_tile = ptrs_sched.load_k(config, k_ptrs, block_sched.n_block_max - 1)
+    # Process n_blocks with seqlen masking
+    for n_block in tl.range(
+        block_sched.n_block_max - 1, block_sched.n_block_max_no_mask - 1, -1
+    ):
+        k_tile, acc_o, row_max, row_sum = _dec_inner_dense_kernel(
+            config=config,
+            ptrs_sched=ptrs_sched,
+            mask_sched=mask_sched,
+            softmax_sched=softmax_sched,
+            q_tile=q_tile,
+            k_tile=k_tile,
+            k_ptrs=k_ptrs,
+            v_ptrs=v_ptrs,
+            acc_o=acc_o,
+            row_max=row_max,
+            row_sum=row_sum,
+            n_block=n_block,
+            n_block_min=block_sched.n_block_max_no_mask,
+            IS_MASK=True,
+            MASK_LOCAL=True if IS_LOCAL else False,
+            MASK_SINK=False,
+            CHECK_INF=True,
+        )
+    # Process n_blocks without masking
+    if not IS_LOCAL and block_sched.n_block_max_no_mask > block_sched.n_block_min:
+        # Load key tile
+        k_tile = ptrs_sched.load_k(config, k_ptrs, block_sched.n_block_max_no_mask - 1)
+        for n_block in tl.range(
+            block_sched.n_block_max_no_mask - 1, block_sched.n_block_min - 1, -1
+        ):
+            k_tile, acc_o, row_max, row_sum = _dec_inner_dense_kernel(
+                config=config,
+                ptrs_sched=ptrs_sched,
+                mask_sched=mask_sched,
+                softmax_sched=softmax_sched,
+                q_tile=q_tile,
+                k_tile=k_tile,
+                k_ptrs=k_ptrs,
+                v_ptrs=v_ptrs,
+                acc_o=acc_o,
+                row_max=row_max,
+                row_sum=row_sum,
+                n_block=n_block,
+                n_block_min=block_sched.n_block_min,
+                IS_MASK=False,
+                MASK_LOCAL=False,
+                MASK_SINK=False,
+                CHECK_INF=False,
+            )
+    if IS_LOCAL:
+        # Process n_blocks with local right masking
+        if block_sched.n_block_window_max > block_sched.n_block_window_max_no_mask:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_max - 1
+            )
+            for n_block in tl.range(
+                block_sched.n_block_window_max - 1,
+                block_sched.n_block_window_max_no_mask - 1,
+                -1,
+            ):
+                k_tile, acc_o, row_max, row_sum = _dec_inner_dense_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    q_tile=q_tile,
+                    k_tile=k_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    acc_o=acc_o,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_max_no_mask,
+                    IS_MASK=True,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                    CHECK_INF=True,
+                )
+        # Process n_blocks without masking
+        if (
+            block_sched.n_block_window_max_no_mask
+            > block_sched.n_block_window_min_no_mask
+        ):
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_max_no_mask - 1
+            )
+            for n_block in tl.range(
+                block_sched.n_block_window_max_no_mask - 1,
+                block_sched.n_block_window_min_no_mask - 1,
+                -1,
+            ):
+                k_tile, acc_o, row_max, row_sum = _dec_inner_dense_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    q_tile=q_tile,
+                    k_tile=k_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    acc_o=acc_o,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_min_no_mask,
+                    IS_MASK=False,
+                    MASK_LOCAL=False,
+                    MASK_SINK=False,
+                    CHECK_INF=False,
+                )
+        # Process n_blocks with local left masking
+        if block_sched.n_block_window_min_no_mask > block_sched.n_block_window_min:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_min_no_mask - 1
+            )
+            for n_block in tl.range(
+                block_sched.n_block_window_min_no_mask - 1,
+                block_sched.n_block_window_min - 1,
+                -1,
+            ):
+                k_tile, acc_o, row_max, row_sum = _dec_inner_dense_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    q_tile=q_tile,
+                    k_tile=k_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    acc_o=acc_o,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_min,
+                    IS_MASK=True,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                    CHECK_INF=True,
+                )
+        # Process n_blocks with local sink masking
+        if block_sched.n_block_sink_max > block_sched.n_block_sink_min:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(config, k_ptrs, block_sched.n_block_sink_max - 1)
+            for n_block in tl.range(
+                block_sched.n_block_sink_max - 1,
+                block_sched.n_block_sink_min - 1,
+                -1,
+            ):
+                k_tile, acc_o, row_max, row_sum = _dec_inner_dense_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    q_tile=q_tile,
+                    k_tile=k_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    acc_o=acc_o,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_sink_min,
+                    IS_MASK=True,
+                    MASK_LOCAL=True,
+                    MASK_SINK=True,
+                    CHECK_INF=True,
+                )
+    # Finalize softmax
+    row_scale, lse_tile = softmax_sched.finalize(
+        row_max=row_max,
+        row_sum=row_sum,
+        IS_LOG2=True,
+    )
+    # Store LSE
+    ptrs_sched.store_lse(config, lse_ptrs, lse_tile)
+    # Finalize rescale
+    acc_o = softmax_sched.rescale_o(
+        acc_o=acc_o,
+        row_scale=row_scale,
+    )
+    # Store output
+    ptrs_sched.store_out(config, out_ptrs, acc_o)
+_dec_dense_kernel = cache_utils.wrap_kernel(_dec_dense_kernel)
+_dec_dense_autotuned_kernel = None
+def _get_autotuned_kernel():
+    global _dec_dense_autotuned_kernel
+    if _dec_dense_autotuned_kernel is None:
+        jit_kernel = _dec_dense_kernel._kernel
+        autotuned = autotuner.make_dec_dense_autotuned_kernel(jit_kernel)
+        _dec_dense_autotuned_kernel = autotuner.AutotunedKernel(autotuned)
+    return _dec_dense_autotuned_kernel
+def _flash_dense_attn_decode(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    softmax_scale: float = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    batch_size, num_heads_q, head_dim = query.shape
+    _, seqlen_k, num_heads_kv, _ = key.shape
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_dec_inputs(
+            query,
+            key,
+            value,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            cu_seqlens_k=None,
+            seqused_k=None,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name="dec_dense",
+        seqlen_q=1,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _dec_dense_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        TILE_M = max(triton.next_power_of_2(qhead_per_kvhead), 16)
+        TILE_N = 128
+        num_warps = num_stages = num_ctas = None
+    num_splits = utils.num_splits_heuristic(
+        seqlen_q=qhead_per_kvhead,
+        seqlen_k=seqlen_k,
+        num_SMs=num_SMs,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        max_split_blocks=utils.max_split_blocks_from_window_sizes(window_sizes, TILE_N)
+        if is_local
+        else None,
+    )
+    out_dtype = torch.bfloat16 if is_quant else query.dtype
+    out = (
+        out
+        if out is not None
+        else torch.empty(query.shape, dtype=out_dtype, device=device)
+    )
+    lse = (
+        lse
+        if lse is not None
+        else torch.empty((batch_size, num_heads_q), dtype=torch.float32, device=device)
+    )
+    out_partial = torch.empty(
+        (num_splits, batch_size, num_heads_q, head_dim),
+        dtype=torch.float32,
+        device=device,
+    )
+    lse_partial = torch.empty(
+        (num_splits, batch_size, num_heads_q),
+        dtype=torch.float32,
+        device=device,
+    )
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    grid = launch_grid.get_dec_grid(
+        batch_size=batch_size,
+        num_heads_kv=num_heads_kv,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        out_partial,
+        lse_partial,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        query.stride(0),
+        query.stride(-2),
+        1,
+        key.stride(0),
+        key.stride(-2),
+        key.stride(-3),
+        value.stride(0),
+        value.stride(-2),
+        value.stride(-3),
+        out_partial.stride(1),
+        out_partial.stride(-2),
+        1,
+        out_partial.stride(0),
+        lse_partial.stride(1),
+        lse_partial.stride(-1),
+        1,
+        lse_partial.stride(0),
+        window_sizes.stride(0),
+        None,
+        None,
+        None,
+        None,
+        num_splits,
+        seqlen_q=1,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=0,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD_PACKGQA=qhead_per_kvhead,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_LOCAL=is_local,
+        HAS_CU_SEQLENS_Q=False,
+        HAS_CU_SEQLENS_K=False,
+        HAS_SEQUSED_Q=False,
+        HAS_SEQUSED_K=False,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name="dec_dense",
+                seqlen_q=1,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_quant=is_quant,
+            )
+    flash_dec_combine._flash_attn_dec_combine(
+        out_partial,
+        lse_partial,
+        out,
+        lse,
+    )
+    return out, lse
+def _flash_dense_attn_varlen_decode(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_k: int,
+    softmax_scale: float = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    seqused_k: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    batch_size, num_heads_q, head_dim = query.shape
+    _, num_heads_kv, _ = key.shape
+    seqlen_k = max_seqlen_k
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_dec_inputs(
+            query,
+            key,
+            value,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            cu_seqlens_k=cu_seqlens_k,
+            seqused_k=seqused_k,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name="dec_dense",
+        seqlen_q=1,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _dec_dense_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        TILE_M = max(triton.next_power_of_2(qhead_per_kvhead), 16)
+        TILE_N = 128
+        num_warps = num_stages = num_ctas = None
+    num_splits = utils.num_splits_heuristic(
+        seqlen_q=qhead_per_kvhead,
+        seqlen_k=seqlen_k,
+        num_SMs=num_SMs,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        max_split_blocks=utils.max_split_blocks_from_window_sizes(window_sizes, TILE_N)
+        if is_local
+        else None,
+    )
+    out_dtype = torch.bfloat16 if is_quant else query.dtype
+    out = (
+        out
+        if out is not None
+        else torch.empty(query.shape, dtype=out_dtype, device=device)
+    )
+    lse = (
+        lse
+        if lse is not None
+        else torch.empty((batch_size, num_heads_q), dtype=torch.float32, device=device)
+    )
+    out_partial = torch.empty(
+        (num_splits, batch_size, num_heads_q, head_dim),
+        dtype=torch.float32,
+        device=device,
+    )
+    lse_partial = torch.empty(
+        (num_splits, batch_size, num_heads_q),
+        dtype=torch.float32,
+        device=device,
+    )
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    grid = launch_grid.get_dec_grid(
+        batch_size=batch_size,
+        num_heads_kv=num_heads_kv,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        out_partial,
+        lse_partial,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        query.stride(0),
+        query.stride(-2),
+        1,
+        0,
+        key.stride(-2),
+        key.stride(0),
+        0,
+        value.stride(-2),
+        value.stride(0),
+        out_partial.stride(1),
+        out_partial.stride(-2),
+        1,
+        out_partial.stride(0),
+        lse_partial.stride(1),
+        lse_partial.stride(-1),
+        1,
+        lse_partial.stride(0),
+        window_sizes.stride(0),
+        None,
+        cu_seqlens_k,
+        None,
+        seqused_k,
+        num_splits,
+        seqlen_q=1,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=0,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD_PACKGQA=qhead_per_kvhead,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_LOCAL=is_local,
+        HAS_CU_SEQLENS_Q=False,
+        HAS_CU_SEQLENS_K=True,
+        HAS_SEQUSED_Q=False,
+        HAS_SEQUSED_K=seqused_k is not None,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name="dec_dense",
+                seqlen_q=1,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_quant=is_quant,
+            )
+    flash_dec_combine._flash_attn_dec_combine(
+        out_partial,
+        lse_partial,
+        out,
+        lse,
+    )
+    return out, lse

build/torch-cuda/flash_dense_fwd.py ADDED Viewed

	@@ -0,0 +1,957 @@

+from typing import Tuple, Optional
+import torch
+import triton
+import triton.language as tl
+from . import (
+    assert_inputs,
+    utils,
+    cache_utils,
+    launch_template,
+    launch_grid,
+    flash_fwd_combine,
+    kernel_repr,
+    autotuner,
+)
+from .scheduler import (
+    AttnFwdGridIndex,
+    AttnFwdConfig,
+    AttnFwdBlockScheduler,
+    AttnFwdPointerScheduler,
+    AttnMaskScheduler,
+    SoftmaxScheduler,
+)
+@triton.jit
+def _fwd_inner_dense_kernel(
+    config: AttnFwdConfig,
+    ptrs_sched: AttnFwdPointerScheduler,
+    mask_sched: AttnMaskScheduler,
+    softmax_sched: SoftmaxScheduler,
+    q_tile,
+    k_tile,
+    k_ptrs,
+    v_ptrs,
+    acc_o,
+    row_max,
+    row_sum,
+    n_block,
+    n_block_min,
+    IS_MASK: tl.constexpr,
+    MASK_CAUSAL: tl.constexpr,
+    MASK_LOCAL: tl.constexpr,
+    MASK_SINK: tl.constexpr,
+    CHECK_INF: tl.constexpr,
+):
+    # Compute attention scores
+    acc_s = tl.dot(q_tile, k_tile.T)
+    # Prefetch next key tile
+    if n_block > n_block_min:
+        # Load next key tile
+        k_tile = ptrs_sched.load_k(config, k_ptrs, n_block - 1)
+    if IS_MASK:
+        # Apply mask to attention scores
+        acc_s = mask_sched.apply_mask(
+            acc_s=acc_s,
+            iter_block=n_block,
+            MASK_CAUSAL=MASK_CAUSAL,
+            MASK_LOCAL=MASK_LOCAL,
+            MASK_SINK=MASK_SINK,
+        )
+    # Apply online softmax
+    p, row_max, row_sum, row_scale = softmax_sched.online_softmax(
+        acc_s=acc_s,
+        row_max=row_max,
+        row_sum=row_sum,
+        CHECK_INF=CHECK_INF,
+    )
+    # Load value tile
+    v_tile = ptrs_sched.load_v(config, v_ptrs, n_block)
+    # Rescale output accumulator
+    acc_o = softmax_sched.rescale_o(
+        acc_o=acc_o,
+        row_scale=row_scale,
+    )
+    # Update output accumulator
+    acc_o += tl.dot(p.to(v_tile.dtype), v_tile)
+    return k_tile, acc_o, row_max, row_sum
+@triton.jit(repr=kernel_repr.fwd_dense_repr)
+def _fwd_dense_kernel(
+    Q,
+    K,
+    V,
+    Out,
+    Lse,
+    softmax_scale,
+    query_scale,
+    key_scale,
+    value_scale,
+    window_sizes,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_os,
+    stride_lb,
+    stride_lh,
+    stride_ls,
+    stride_wh,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    seqused_q,
+    seqused_k,
+    num_splits,
+    seqlen_q,
+    seqlen_k,
+    head_dim,
+    SEQLEN_Q_CACHE: tl.constexpr,
+    SEQLEN_K_CACHE: tl.constexpr,
+    QHEAD_PER_KVHEAD: tl.constexpr,
+    PACK_GQA: tl.constexpr,
+    QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+    TILE_K: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    IS_LOCAL: tl.constexpr,
+    IS_SPLIT_KV: tl.constexpr,
+    HAS_CU_SEQLENS_Q: tl.constexpr,
+    HAS_CU_SEQLENS_K: tl.constexpr,
+    HAS_SEQUSED_Q: tl.constexpr,
+    HAS_SEQUSED_K: tl.constexpr,
+):
+    # Create grid index
+    grid_idx = AttnFwdGridIndex.create(
+        num_splits=num_splits,
+        QHEAD_PER_KVHEAD=QHEAD_PER_KVHEAD,
+        IS_SPLIT_KV=IS_SPLIT_KV,
+        PACK_GQA=PACK_GQA,
+    )
+    # Load window sizes
+    (
+        window_size_sink,
+        window_size_left,
+        window_size_right,
+        window_size_dist,
+    ) = grid_idx.load_window_sizes(
+        window_sizes=window_sizes,
+        stride_wh=stride_wh,
+        IS_LOCAL=IS_LOCAL,
+    )
+    # Create config
+    config = AttnFwdConfig.create(
+        softmax_scale=softmax_scale,
+        query_scale=query_scale,
+        key_scale=key_scale,
+        value_scale=value_scale,
+        m_block=grid_idx.m_block,
+        batch_idx=grid_idx.batch_idx,
+        window_size_sink=window_size_sink,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        window_size_dist=window_size_dist,
+        head_dim=head_dim,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        PACK_GQA=PACK_GQA,
+        QHEAD_PER_KVHEAD_PACKGQA=QHEAD_PER_KVHEAD_PACKGQA,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+        HAS_SEQUSED_Q=HAS_SEQUSED_Q,
+        HAS_SEQUSED_K=HAS_SEQUSED_K,
+    )
+    # Create pointer scheduler
+    ptrs_sched = AttnFwdPointerScheduler.create(
+        config=config,
+        Q=Q,
+        K=K,
+        V=V,
+        Out=Out,
+        Lse=Lse,
+        batch_idx=grid_idx.batch_idx,
+        head_idx=grid_idx.head_idx,
+        head_kv_idx=grid_idx.head_kv_idx,
+        split_idx=grid_idx.split_idx,
+        stride_qb=stride_qb,
+        stride_qh=stride_qh,
+        stride_qm=stride_qm,
+        stride_kb=stride_kb,
+        stride_kh=stride_kh,
+        stride_kn=stride_kn,
+        stride_vb=stride_vb,
+        stride_vh=stride_vh,
+        stride_vn=stride_vn,
+        stride_ob=stride_ob,
+        stride_oh=stride_oh,
+        stride_om=stride_om,
+        stride_os=stride_os,
+        stride_lb=stride_lb,
+        stride_lh=stride_lh,
+        stride_ls=stride_ls,
+        IS_SPLIT_KV=IS_SPLIT_KV,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+    )
+    # Create block scheduler
+    block_sched = AttnFwdBlockScheduler.create(
+        config=config,
+        split_idx=grid_idx.split_idx,
+        num_splits=num_splits,
+        IS_CAUSAL=IS_CAUSAL,
+        IS_LOCAL=IS_LOCAL,
+        IS_SPLIT_KV=IS_SPLIT_KV,
+    )
+    # Create mask scheduler
+    mask_sched = AttnMaskScheduler.create(config)
+    # Create softmax scheduler
+    softmax_sched = SoftmaxScheduler.create(config)
+    # Create pointers
+    out_ptrs = ptrs_sched.make_out_ptrs(config)
+    lse_ptrs = ptrs_sched.make_lse_ptrs(config)
+    # Early exit if no n_blocks to process
+    if block_sched.is_empty():
+        ptrs_sched.store_empty(config, out_ptrs, lse_ptrs, IS_SPLIT_KV=IS_SPLIT_KV)
+        return
+    q_ptrs = ptrs_sched.make_q_ptrs(config)
+    k_ptrs = ptrs_sched.make_k_ptrs(config)
+    v_ptrs = ptrs_sched.make_v_ptrs(config)
+    # Initialize accumulators
+    row_max = tl.full((TILE_M,), float("-inf"), dtype=tl.float32)
+    row_sum = tl.zeros((TILE_M,), dtype=tl.float32)
+    acc_o = tl.zeros((TILE_M, TILE_K), dtype=tl.float32)
+    # Load query tile
+    q_tile = ptrs_sched.load_q(config, q_ptrs)
+    # Load key tile
+    k_tile = ptrs_sched.load_k(config, k_ptrs, block_sched.n_block_max - 1)
+    # Process n_blocks with causal masking
+    if IS_CAUSAL or IS_LOCAL:
+        n_block_max_no_mask = block_sched.n_block_max_no_mask
+        for n_block in tl.range(
+            block_sched.n_block_max - 1, block_sched.n_block_max_no_mask - 1, -1
+        ):
+            k_tile, acc_o, row_max, row_sum = _fwd_inner_dense_kernel(
+                config=config,
+                ptrs_sched=ptrs_sched,
+                mask_sched=mask_sched,
+                softmax_sched=softmax_sched,
+                q_tile=q_tile,
+                k_tile=k_tile,
+                k_ptrs=k_ptrs,
+                v_ptrs=v_ptrs,
+                acc_o=acc_o,
+                row_max=row_max,
+                row_sum=row_sum,
+                n_block=n_block,
+                n_block_min=block_sched.n_block_max_no_mask,
+                IS_MASK=True,
+                MASK_CAUSAL=IS_CAUSAL,
+                MASK_LOCAL=True if IS_LOCAL else False,
+                MASK_SINK=False,
+                CHECK_INF=True,
+            )
+    else:
+        # First iteration with seqlen masking
+        n_block = block_sched.n_block_max - 1
+        # Triton aggregate does not support attribute reassignment yet
+        # block_sched.n_block_max_no_mask = n_block
+        n_block_max_no_mask = n_block
+        k_tile, acc_o, row_max, row_sum = _fwd_inner_dense_kernel(
+            config=config,
+            ptrs_sched=ptrs_sched,
+            mask_sched=mask_sched,
+            softmax_sched=softmax_sched,
+            q_tile=q_tile,
+            k_tile=k_tile,
+            k_ptrs=k_ptrs,
+            v_ptrs=v_ptrs,
+            acc_o=acc_o,
+            row_max=row_max,
+            row_sum=row_sum,
+            n_block=n_block,
+            n_block_min=n_block,
+            IS_MASK=True,
+            MASK_CAUSAL=False,
+            MASK_LOCAL=False,
+            MASK_SINK=False,
+            CHECK_INF=True,
+        )
+    # Process n_blocks without masking
+    if not IS_LOCAL and n_block_max_no_mask > block_sched.n_block_min:
+        # Load key tile
+        k_tile = ptrs_sched.load_k(config, k_ptrs, n_block_max_no_mask - 1)
+        for n_block in tl.range(
+            n_block_max_no_mask - 1, block_sched.n_block_min - 1, -1
+        ):
+            k_tile, acc_o, row_max, row_sum = _fwd_inner_dense_kernel(
+                config=config,
+                ptrs_sched=ptrs_sched,
+                mask_sched=mask_sched,
+                softmax_sched=softmax_sched,
+                q_tile=q_tile,
+                k_tile=k_tile,
+                k_ptrs=k_ptrs,
+                v_ptrs=v_ptrs,
+                acc_o=acc_o,
+                row_max=row_max,
+                row_sum=row_sum,
+                n_block=n_block,
+                n_block_min=block_sched.n_block_min,
+                IS_MASK=False,
+                MASK_CAUSAL=False,
+                MASK_LOCAL=False,
+                MASK_SINK=False,
+                CHECK_INF=False,
+            )
+    if IS_LOCAL:
+        # Process n_blocks with local right masking
+        if block_sched.n_block_window_max > block_sched.n_block_window_max_no_mask:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_max - 1
+            )
+            for n_block in tl.range(
+                block_sched.n_block_window_max - 1,
+                block_sched.n_block_window_max_no_mask - 1,
+                -1,
+            ):
+                k_tile, acc_o, row_max, row_sum = _fwd_inner_dense_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    q_tile=q_tile,
+                    k_tile=k_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    acc_o=acc_o,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_max_no_mask,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                    CHECK_INF=True,
+                )
+        # Process n_blocks without masking
+        if (
+            block_sched.n_block_window_max_no_mask
+            > block_sched.n_block_window_min_no_mask
+        ):
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_max_no_mask - 1
+            )
+            for n_block in tl.range(
+                block_sched.n_block_window_max_no_mask - 1,
+                block_sched.n_block_window_min_no_mask - 1,
+                -1,
+            ):
+                k_tile, acc_o, row_max, row_sum = _fwd_inner_dense_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    q_tile=q_tile,
+                    k_tile=k_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    acc_o=acc_o,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_min_no_mask,
+                    IS_MASK=False,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=False,
+                    MASK_SINK=False,
+                    CHECK_INF=False,
+                )
+        # Process n_blocks with local left masking
+        if block_sched.n_block_window_min_no_mask > block_sched.n_block_window_min:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_min_no_mask - 1
+            )
+            for n_block in tl.range(
+                block_sched.n_block_window_min_no_mask - 1,
+                block_sched.n_block_window_min - 1,
+                -1,
+            ):
+                k_tile, acc_o, row_max, row_sum = _fwd_inner_dense_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    q_tile=q_tile,
+                    k_tile=k_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    acc_o=acc_o,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_min,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                    CHECK_INF=True,
+                )
+        # Process n_blocks with local sink masking
+        if block_sched.n_block_sink_max > block_sched.n_block_sink_min:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(config, k_ptrs, block_sched.n_block_sink_max - 1)
+            for n_block in tl.range(
+                block_sched.n_block_sink_max - 1,
+                block_sched.n_block_sink_min - 1,
+                -1,
+            ):
+                k_tile, acc_o, row_max, row_sum = _fwd_inner_dense_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    q_tile=q_tile,
+                    k_tile=k_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    acc_o=acc_o,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_sink_min,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=True,
+                    CHECK_INF=True,
+                )
+    # Finalize softmax
+    row_scale, lse_tile = softmax_sched.finalize(
+        row_max=row_max,
+        row_sum=row_sum,
+        IS_LOG2=IS_SPLIT_KV,
+    )
+    # Store LSE
+    ptrs_sched.store_lse(config, lse_ptrs, lse_tile)
+    # Finalize rescale
+    acc_o = softmax_sched.rescale_o(
+        acc_o=acc_o,
+        row_scale=row_scale,
+    )
+    # Store output
+    ptrs_sched.store_out(config, out_ptrs, acc_o, IS_SPLIT_KV=IS_SPLIT_KV)
+_fwd_dense_kernel = cache_utils.wrap_kernel(_fwd_dense_kernel)
+_fwd_dense_kernel_autotuned = None
+def _get_autotuned_kernel():
+    global _fwd_dense_kernel_autotuned
+    if _fwd_dense_kernel_autotuned is None:
+        jit_kernel = _fwd_dense_kernel._kernel
+        autotuned = autotuner.make_fwd_dense_autotuned_kernel(jit_kernel)
+        _fwd_dense_kernel_autotuned = autotuner.AutotunedKernel(autotuned)
+    return _fwd_dense_kernel_autotuned
+def _flash_dense_attn_forward(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    is_causal: bool = False,
+    softmax_scale: float = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    is_split_kv: bool = False,
+    pack_gqa: bool = False,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, float]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    batch_size, seqlen_q, num_heads_q, head_dim = query.shape
+    _, seqlen_k, num_heads_kv, _ = key.shape
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    qhead_per_kvhead_packgqa = num_heads_q // num_heads_kv if pack_gqa else 1
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_fwd_inputs(
+            query,
+            key,
+            value,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            cu_seqlens_q=None,
+            cu_seqlens_k=None,
+            seqused_q=None,
+            seqused_k=None,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    _kernel_name = f"fwd_dense{'_split' if is_split_kv else ''}"
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name=_kernel_name,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_causal=is_causal,
+        pack_gqa=pack_gqa,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _fwd_dense_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        # Placeholder for pre-launch computations
+        TILE_M = TILE_N = 64
+        num_warps = num_stages = num_ctas = None
+    num_splits = (
+        utils.num_splits_heuristic(
+            seqlen_q=seqlen_q * qhead_per_kvhead_packgqa,
+            seqlen_k=seqlen_k,
+            num_SMs=num_SMs,
+            TILE_M=TILE_M,
+            TILE_N=TILE_N,
+            max_split_blocks=utils.max_split_blocks_from_window_sizes(
+                window_sizes, TILE_N
+            )
+            if is_local
+            else None,
+        )
+        if is_split_kv
+        else 1
+    )
+    out_dtype = torch.bfloat16 if is_quant else query.dtype
+    out = out if out is not None else torch.empty_like(query, dtype=out_dtype)
+    lse = (
+        lse
+        if lse is not None
+        else torch.empty(
+            (batch_size, num_heads_q, seqlen_q),
+            dtype=torch.float32,
+            device=query.device,
+        )
+    )
+    if is_split_kv:
+        out_partial = torch.empty(
+            (num_splits, batch_size, seqlen_q, num_heads_q, head_dim),
+            dtype=torch.float32,
+            device=query.device,
+        )
+        lse_partial = torch.empty(
+            (num_splits, batch_size, num_heads_q, seqlen_q),
+            dtype=torch.float32,
+            device=query.device,
+        )
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    grid = launch_grid.get_fwd_grid(
+        batch_size=batch_size,
+        seqlen_q=seqlen_q,
+        num_heads_q=num_heads_q,
+        num_heads_kv=num_heads_kv,
+        pack_gqa=pack_gqa,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        out if not is_split_kv else out_partial,
+        lse if not is_split_kv else lse_partial,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        query.stride(0),
+        query.stride(-2),
+        query.stride(-3),
+        key.stride(0),
+        key.stride(-2),
+        key.stride(-3),
+        value.stride(0),
+        value.stride(-2),
+        value.stride(-3),
+        out.stride(0) if not is_split_kv else out_partial.stride(1),
+        out.stride(-2) if not is_split_kv else out_partial.stride(-2),
+        out.stride(-3) if not is_split_kv else out_partial.stride(-3),
+        0 if not is_split_kv else out_partial.stride(0),
+        lse.stride(0) if not is_split_kv else lse_partial.stride(1),
+        lse.stride(-2) if not is_split_kv else lse_partial.stride(-2),
+        0 if not is_split_kv else lse_partial.stride(0),
+        window_sizes.stride(0),
+        None,
+        None,
+        None,
+        None,
+        num_splits,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=seqlen_q // 1024,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD=qhead_per_kvhead,
+        PACK_GQA=pack_gqa,
+        QHEAD_PER_KVHEAD_PACKGQA=qhead_per_kvhead_packgqa,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_CAUSAL=is_causal,
+        IS_LOCAL=is_local,
+        IS_SPLIT_KV=is_split_kv,
+        HAS_CU_SEQLENS_Q=False,
+        HAS_CU_SEQLENS_K=False,
+        HAS_SEQUSED_Q=False,
+        HAS_SEQUSED_K=False,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name=_kernel_name,
+                seqlen_q=seqlen_q,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_causal=is_causal,
+                pack_gqa=pack_gqa,
+                is_quant=is_quant,
+            )
+    if is_split_kv:
+        flash_fwd_combine._flash_attn_fwd_combine(
+            out_partial,
+            lse_partial,
+            out,
+            lse,
+        )
+    return out, lse, softmax_scale
+def _flash_dense_attn_varlen_forward(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    is_causal: bool = False,
+    softmax_scale: float = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    is_split_kv: bool = False,
+    pack_gqa: bool = False,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, float]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    total_seqlen_q, num_heads_q, head_dim = query.shape
+    _, num_heads_kv, _ = key.shape
+    batch_size = cu_seqlens_q.shape[0] - 1
+    seqlen_q = max_seqlen_q
+    seqlen_k = max_seqlen_k
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    qhead_per_kvhead_packgqa = num_heads_q // num_heads_kv if pack_gqa else 1
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_fwd_inputs(
+            query,
+            key,
+            value,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    _kernel_name = f"fwd_dense{'_split' if is_split_kv else ''}"
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name=_kernel_name,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_causal=is_causal,
+        pack_gqa=pack_gqa,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _fwd_dense_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        # Placeholder for pre-launch computations
+        TILE_M = TILE_N = 64
+        num_warps = num_stages = num_ctas = None
+    num_splits = (
+        utils.num_splits_heuristic(
+            seqlen_q=seqlen_q * qhead_per_kvhead_packgqa,
+            seqlen_k=seqlen_k,
+            num_SMs=num_SMs,
+            TILE_M=TILE_M,
+            TILE_N=TILE_N,
+            max_split_blocks=utils.max_split_blocks_from_window_sizes(
+                window_sizes, TILE_N
+            )
+            if is_local
+            else None,
+        )
+        if is_split_kv
+        else 1
+    )
+    out_dtype = torch.bfloat16 if is_quant else query.dtype
+    out = torch.empty_like(query, dtype=out_dtype)
+    lse = torch.empty(
+        (num_heads_q, total_seqlen_q),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    if is_split_kv:
+        out_partial = torch.empty(
+            (num_splits, total_seqlen_q, num_heads_q, head_dim),
+            dtype=torch.float32,
+            device=query.device,
+        )
+        lse_partial = torch.empty(
+            (num_splits, num_heads_q, total_seqlen_q),
+            dtype=torch.float32,
+            device=query.device,
+        )
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    grid = launch_grid.get_fwd_grid(
+        batch_size=batch_size,
+        seqlen_q=seqlen_q,
+        num_heads_q=num_heads_q,
+        num_heads_kv=num_heads_kv,
+        pack_gqa=pack_gqa,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        out if not is_split_kv else out_partial,
+        lse if not is_split_kv else lse_partial,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        0,
+        query.stride(-2),
+        query.stride(0),
+        0,
+        key.stride(-2),
+        key.stride(0),
+        0,
+        value.stride(-2),
+        value.stride(0),
+        0,
+        out.stride(-2) if not is_split_kv else out_partial.stride(-2),
+        out.stride(0) if not is_split_kv else out_partial.stride(-3),
+        0 if not is_split_kv else out_partial.stride(0),
+        0,
+        lse.stride(-2) if not is_split_kv else lse_partial.stride(-2),
+        0 if not is_split_kv else lse_partial.stride(0),
+        window_sizes.stride(0),
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_q,
+        seqused_k,
+        num_splits,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=seqlen_q // 1024,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD=qhead_per_kvhead,
+        PACK_GQA=pack_gqa,
+        QHEAD_PER_KVHEAD_PACKGQA=qhead_per_kvhead_packgqa,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_CAUSAL=is_causal,
+        IS_LOCAL=is_local,
+        IS_SPLIT_KV=is_split_kv,
+        HAS_CU_SEQLENS_Q=True,
+        HAS_CU_SEQLENS_K=True,
+        HAS_SEQUSED_Q=seqused_q is not None,
+        HAS_SEQUSED_K=seqused_k is not None,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name=_kernel_name,
+                seqlen_q=seqlen_q,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_causal=is_causal,
+                pack_gqa=pack_gqa,
+                is_quant=is_quant,
+            )
+    if is_split_kv:
+        flash_fwd_combine._flash_attn_fwd_combine(
+            out_partial,
+            lse_partial,
+            out,
+            lse,
+            cu_seqlens_q=cu_seqlens_q,
+            seqused_q=seqused_q,
+        )
+    return out, lse, softmax_scale

build/torch-cuda/flash_fwd_combine.py ADDED Viewed

	@@ -0,0 +1,239 @@

+import torch
+import triton
+import triton.language as tl
+from . import (
+    utils,
+    cache_utils,
+    launch_template,
+    launch_grid,
+    seqlen_info,
+    kernel_repr,
+)
+@triton.jit(repr=kernel_repr.fwd_combine_repr)
+def _fwd_combine_kernel(
+    Out_partial,
+    Lse_partial,
+    Out,
+    Lse,
+    stride_ops,
+    stride_opb,
+    stride_oph,
+    stride_opm,
+    stride_lps,
+    stride_lpb,
+    stride_lph,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_lb,
+    stride_lh,
+    cu_seqlens_q,
+    seqused_q,
+    num_splits,
+    seqlen_q,
+    num_heads_q,
+    head_dim,
+    TILE_M: tl.constexpr,
+    TILE_K: tl.constexpr,
+    HAS_CU_SEQLENS_Q: tl.constexpr,
+    HAS_SEQUSED_Q: tl.constexpr,
+):
+    m_block = tl.program_id(0)
+    bh_idx = tl.program_id(1)
+    batch_idx = bh_idx // num_heads_q
+    head_idx = bh_idx - batch_idx * num_heads_q
+    # Get seqlen info for this batch
+    offset_q, actual_seqlen_q = seqlen_info.get_seqlen_info(
+        batch_idx=batch_idx,
+        seqlen_static=seqlen_q,
+        cu_seqlens=cu_seqlens_q,
+        seqused=seqused_q,
+        HAS_CU_SEQLENS=HAS_CU_SEQLENS_Q,
+        HAS_SEQUSED=HAS_SEQUSED_Q,
+    )
+    # Initialize base pointers
+    out_part_base = seqlen_info.offset_batch_Q(
+        Out_partial + head_idx * stride_oph,
+        batch_idx,
+        offset_q,
+        0,
+        stride_opb,
+        stride_opm,
+        HAS_CU_SEQLENS_Q,
+        USE_PADDED=False,
+    )
+    lse_part_base = seqlen_info.offset_batch_Q(
+        Lse_partial + head_idx * stride_lph,
+        batch_idx,
+        offset_q,
+        0,
+        stride_lpb,
+        1,
+        HAS_CU_SEQLENS_Q,
+        USE_PADDED=False,
+    )
+    out_base = seqlen_info.offset_batch_Q(
+        Out + head_idx * stride_oh,
+        batch_idx,
+        offset_q,
+        0,
+        stride_ob,
+        stride_om,
+        HAS_CU_SEQLENS_Q,
+        USE_PADDED=False,
+    )
+    lse_base = seqlen_info.offset_batch_Q(
+        Lse + head_idx * stride_lh,
+        batch_idx,
+        offset_q,
+        0,
+        stride_lb,
+        1,
+        HAS_CU_SEQLENS_Q,
+        USE_PADDED=False,
+    )
+    # Create descriptors
+    out_part_desc = tl.make_tensor_descriptor(
+        base=out_part_base,
+        shape=[num_splits, actual_seqlen_q, head_dim],
+        strides=[stride_ops, stride_opm, 1],
+        block_shape=[1, TILE_M, TILE_K],
+    )
+    lse_part_desc = tl.make_tensor_descriptor(
+        base=lse_part_base,
+        shape=[num_splits, actual_seqlen_q],
+        strides=[stride_lps, 1],
+        block_shape=[1, TILE_M],
+    )
+    out_desc = tl.make_tensor_descriptor(
+        base=out_base,
+        shape=[actual_seqlen_q, head_dim],
+        strides=[stride_om, 1],
+        block_shape=[TILE_M, TILE_K],
+    )
+    lse_desc = tl.make_tensor_descriptor(
+        base=lse_base,
+        shape=[actual_seqlen_q],
+        strides=[1],
+        block_shape=[TILE_M],
+    )
+    # Initialize accumulators
+    e_sum = tl.zeros((TILE_M,), dtype=tl.float32)
+    e_max = tl.full((TILE_M,), float("-inf"), dtype=tl.float32)
+    acc_o = tl.zeros((TILE_M, TILE_K), dtype=tl.float32)
+    # Combine split outputs
+    for s in tl.range(0, num_splits):
+        # Load partial LSE
+        lse_s = tl.sum(lse_part_desc.load([s, m_block * TILE_M]), axis=0)
+        # Compute normalized exponentials
+        new_e_max = tl.maximum(lse_s, e_max)
+        old_scale = tl.where(e_max == float("-inf"), 0.0, tl.exp2(e_max - new_e_max))
+        exp_logic = tl.where(lse_s == float("-inf"), 0.0, tl.exp2(lse_s - new_e_max))
+        # Load partial outputs
+        o_s = tl.sum(out_part_desc.load([s, m_block * TILE_M, 0]), axis=0)
+        # Compute scaled outputs
+        acc_o *= old_scale[:, None]
+        acc_o += exp_logic[:, None] * o_s
+        # Update e_sum and e_max
+        e_sum = e_sum * old_scale + exp_logic
+        e_max = new_e_max
+    # Normalize output
+    inv_sum = tl.where((e_sum == 0.0) | (e_sum != e_sum), 0.0, 1.0 / e_sum)
+    acc_o *= inv_sum[:, None]
+    # Store output
+    out_desc.store([m_block * TILE_M, 0], acc_o.to(Out.dtype.element_ty))
+    # Compute LSE
+    # ln2 = math.log(2.0)
+    ln2 = 0.6931471805599453
+    lse = tl.where(e_sum > 0.0, (e_max + tl.log2(e_sum)) * ln2, float("-inf"))
+    # Store LSE
+    lse_desc.store([m_block * TILE_M], lse)
+_fwd_combine_kernel = cache_utils.wrap_kernel(_fwd_combine_kernel)
+def _flash_attn_fwd_combine(
+    out_partial: torch.Tensor,
+    lse_partial: torch.Tensor,
+    out: torch.Tensor,
+    lse: torch.Tensor,
+    cu_seqlens_q: torch.Tensor = None,
+    seqused_q: torch.Tensor = None,
+):
+    device = out.device
+    arch = cache_utils.get_device_arch(device)
+    is_varlen = cu_seqlens_q is not None
+    num_splits = out_partial.shape[0]
+    if not is_varlen:
+        batch_size, seqlen_q, num_heads_q, head_dim = out_partial.shape[1:]
+    else:
+        total_q, num_heads_q, head_dim = out_partial.shape[1:]
+        batch_size = cu_seqlens_q.shape[0] - 1
+        seqlen_q = total_q
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    TILE_M, num_warps, num_stages, num_ctas = (
+        launch_template.get_fwd_combine_launch_config(
+            tile_k=TILE_K,
+            device=device,
+            arch=arch,
+        )
+    )
+    grid = launch_grid.get_fwd_combine_grid(
+        batch_size=batch_size,
+        seqlen_q=seqlen_q,
+        num_heads_q=num_heads_q,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    _fwd_combine_kernel[grid](
+        out_partial,
+        lse_partial,
+        out,
+        lse,
+        out_partial.stride(0),
+        out_partial.stride(1) if not is_varlen else 0,
+        out_partial.stride(-2),
+        out_partial.stride(-3),
+        lse_partial.stride(0),
+        lse_partial.stride(1) if not is_varlen else 0,
+        lse_partial.stride(-2),
+        out.stride(0) if not is_varlen else 0,
+        out.stride(-2),
+        out.stride(-3) if not is_varlen else out.stride(0),
+        lse.stride(0) if not is_varlen else 0,
+        lse.stride(-2) if not is_varlen else lse.stride(0),
+        cu_seqlens_q,
+        seqused_q,
+        num_splits,
+        seqlen_q,
+        num_heads_q,
+        head_dim,
+        TILE_M=TILE_M,
+        TILE_K=TILE_K,
+        HAS_CU_SEQLENS_Q=is_varlen,
+        HAS_SEQUSED_Q=seqused_q is not None,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )

build/torch-cuda/flash_gated_bwd.py ADDED Viewed

	@@ -0,0 +1,1317 @@

+from typing import Optional, Tuple
+import math
+import torch
+import triton
+import triton.language as tl
+from . import (
+    assert_inputs,
+    utils,
+    cache_utils,
+    launch_template,
+    launch_grid,
+    activations,
+    flash_bwd_preprocess,
+    flash_bwd_postprocess,
+    kernel_repr,
+    autotuner,
+)
+from .scheduler import (
+    AttnBwdGridIndex,
+    AttnBwdConfig,
+    AttnBwdBlockScheduler,
+    AttnBwdPointerScheduler,
+    AttnMaskScheduler,
+    SoftmaxScheduler,
+)
+@triton.jit
+def _bwd_inner_gated_kernel(
+    config,
+    ptrs_sched,
+    mask_sched,
+    softmax_sched,
+    acc_dk,
+    acc_dv,
+    acc_dd,
+    k_tile,
+    v_tile,
+    d_tile,
+    q_ptrs,
+    a_ptrs,
+    do_ptrs,
+    dq_accum_ptrs,
+    da_ptrs,
+    lse_ptrs,
+    dpsum_ptrs,
+    d_max,
+    d_min,
+    gate_max,
+    m_block,
+    IS_MASK: tl.constexpr,
+    MASK_CAUSAL: tl.constexpr,
+    MASK_LOCAL: tl.constexpr,
+    MASK_SINK: tl.constexpr,
+):
+    skip_gate = False
+    skip_softmax = False
+    # Load alpha tile
+    a_tile = ptrs_sched.load_a(config, a_ptrs, m_block)
+    a_max = tl.max(a_tile)
+    a_min = tl.min(a_tile)
+    # Compute gate threshold for this m_block
+    gate_threshold_log2 = config.get_gate_threshold_log2(
+        m_block=m_block,
+    )
+    # Check if any gates are active for current tile
+    gate_max, skip_gate = softmax_sched.online_gate(
+        a_max=a_max,
+        a_min=a_min,
+        d_max=d_max,
+        d_min=d_min,
+        gate_max=gate_max,
+        gate_threshold_log2=gate_threshold_log2,
+    )
+    if not skip_gate:
+        # Compute attention gates
+        acc_s = d_tile[:, None] * a_tile[None, :]
+        # Compute scaling factor for gated attention score gradients
+        if config.IS_LOGSIGMOID_GATE:
+            ds_scale = tl.sigmoid(-acc_s)
+        else:
+            ds_scale = 1.0
+        # Load query tile
+        q_tile = ptrs_sched.load_q(config, q_ptrs, m_block)
+        # Compute attention scores
+        if config.IS_LOGSIGMOID_GATE:
+            acc_s = softmax_sched.log_sigmoid(
+                acc_s=acc_s,
+            )
+        acc_s += tl.dot(k_tile, q_tile.T)
+        if IS_MASK:
+            # Apply mask to attention scores
+            acc_s = mask_sched.apply_mask(
+                acc_s=acc_s,
+                iter_block=m_block,
+                MASK_CAUSAL=MASK_CAUSAL,
+                MASK_LOCAL=MASK_LOCAL,
+                MASK_SINK=MASK_SINK,
+            )
+        # Load LSE
+        lse_log2 = ptrs_sched.load_lse(config, lse_ptrs, m_block)
+        # Compute attention weights in log2-domain
+        p_log2 = acc_s * config.softmax_scale_log2 - lse_log2[None, :]
+        # Compute softmax threshold for this m_block
+        softmax_threshold_log2 = config.get_softmax_threshold_log2(
+            m_block=m_block,
+        )
+        # Update skip condition based on threshold
+        skip_softmax = tl.max(p_log2 - softmax_threshold_log2[None, :]) < 0.0
+        if not skip_softmax:
+            # Compute attention weights
+            p = activations.exp2(p_log2).to(q_tile.dtype)
+            # Load output gradients tile
+            do_tile = ptrs_sched.load_do(config, do_ptrs, m_block)
+            # Compute value gradients
+            acc_dv += tl.dot(p, do_tile)
+            # Compute attention weight gradients
+            acc_dp = tl.dot(v_tile, tl.trans(do_tile))
+            # Load dpsum
+            dpsum = ptrs_sched.load_dpsum(config, dpsum_ptrs, m_block)
+            # Compute attention score gradients
+            ds = p * (acc_dp - dpsum[None, :]).to(q_tile.dtype)
+            # Compute query gradients
+            dq = tl.dot(tl.trans(ds), k_tile)
+            # Store query gradients
+            ptrs_sched.store_dq(config, dq_accum_ptrs, m_block, dq)
+            # Compute key gradients
+            acc_dk += tl.dot(ds, q_tile)
+            # Compute alpha gradients
+            da = tl.sum(ds * ds_scale * d_tile[:, None], axis=0)
+            # Store alpha gradients
+            ptrs_sched.store_da(config, da_ptrs, m_block, da)
+            # Compute delta gradients
+            acc_dd += tl.sum(ds * ds_scale * a_tile[None, :], axis=1)
+    return (
+        acc_dk,
+        acc_dv,
+        acc_dd,
+        gate_max,
+    )
+@triton.jit(repr=kernel_repr.bwd_gated_repr)
+def _bwd_gated_kernel(
+    Q,
+    K,
+    V,
+    A,
+    D,
+    dO,
+    LSELog2,
+    dPsum,
+    dQaccum,
+    dK,
+    dV,
+    dA,
+    dD,
+    softmax_scale,
+    query_scale,
+    key_scale,
+    value_scale,
+    window_sizes,
+    softmax_threshold,
+    gate_threshold,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_ab,
+    stride_ah,
+    stride_db,
+    stride_dh,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_lb,
+    stride_lh,
+    stride_pb,
+    stride_ph,
+    stride_dqab,
+    stride_dqah,
+    stride_dqam,
+    stride_dkb,
+    stride_dkh,
+    stride_dkn,
+    stride_dks,
+    stride_dvb,
+    stride_dvh,
+    stride_dvn,
+    stride_dvs,
+    stride_dab,
+    stride_dah,
+    stride_ddb,
+    stride_ddh,
+    stride_dds,
+    stride_wh,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    seqused_q,
+    seqused_k,
+    num_splits,
+    seqlen_q,
+    seqlen_k,
+    head_dim,
+    SEQLEN_Q_CACHE: tl.constexpr,
+    SEQLEN_K_CACHE: tl.constexpr,
+    QHEAD_PER_KVHEAD: tl.constexpr,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+    TILE_K: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    IS_LOCAL: tl.constexpr,
+    IS_SPLIT_QO: tl.constexpr,
+    IS_LOGSIGMOID_GATE: tl.constexpr,
+    IS_ADAPT_GATE: tl.constexpr,
+    HAS_CU_SEQLENS_Q: tl.constexpr,
+    HAS_CU_SEQLENS_K: tl.constexpr,
+    HAS_SEQUSED_Q: tl.constexpr,
+    HAS_SEQUSED_K: tl.constexpr,
+):
+    # Create grid index
+    grid_idx = AttnBwdGridIndex.create(
+        num_splits=num_splits,
+        QHEAD_PER_KVHEAD=QHEAD_PER_KVHEAD,
+        IS_SPLIT_QO=IS_SPLIT_QO,
+    )
+    # Load window sizes
+    (
+        window_size_sink,
+        window_size_left,
+        window_size_right,
+        window_size_dist,
+    ) = grid_idx.load_window_sizes(
+        window_sizes=window_sizes,
+        stride_wh=stride_wh,
+        IS_LOCAL=IS_LOCAL,
+    )
+    # Create config
+    config = AttnBwdConfig.create(
+        softmax_scale=softmax_scale,
+        softmax_threshold=softmax_threshold,
+        gate_threshold=gate_threshold,
+        query_scale=query_scale,
+        key_scale=key_scale,
+        value_scale=value_scale,
+        n_block=grid_idx.n_block,
+        batch_idx=grid_idx.batch_idx,
+        window_size_sink=window_size_sink,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        window_size_dist=window_size_dist,
+        head_dim=head_dim,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        QHEAD_PER_KVHEAD=QHEAD_PER_KVHEAD,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_CAUSAL=IS_CAUSAL,
+        IS_LOGSIGMOID_GATE=IS_LOGSIGMOID_GATE,
+        IS_ADAPT_GATE=IS_ADAPT_GATE,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+        HAS_SEQUSED_Q=HAS_SEQUSED_Q,
+        HAS_SEQUSED_K=HAS_SEQUSED_K,
+    )
+    # Create pointer scheduler
+    ptrs_sched = AttnBwdPointerScheduler.create(
+        config=config,
+        Q=Q,
+        K=K,
+        V=V,
+        A=A,
+        D=D,
+        dO=dO,
+        LSELog2=LSELog2,
+        dPsum=dPsum,
+        dQaccum=dQaccum,
+        dK=dK,
+        dV=dV,
+        dA=dA,
+        dD=dD,
+        batch_idx=grid_idx.batch_idx,
+        head_idx=grid_idx.head_idx,
+        head_kv_idx=grid_idx.head_kv_idx,
+        split_idx=grid_idx.split_idx,
+        stride_qb=stride_qb,
+        stride_qh=stride_qh,
+        stride_qm=stride_qm,
+        stride_kb=stride_kb,
+        stride_kh=stride_kh,
+        stride_kn=stride_kn,
+        stride_vb=stride_vb,
+        stride_vh=stride_vh,
+        stride_vn=stride_vn,
+        stride_ab=stride_ab,
+        stride_ah=stride_ah,
+        stride_db=stride_db,
+        stride_dh=stride_dh,
+        stride_dob=stride_dob,
+        stride_doh=stride_doh,
+        stride_dom=stride_dom,
+        stride_lb=stride_lb,
+        stride_lh=stride_lh,
+        stride_pb=stride_pb,
+        stride_ph=stride_ph,
+        stride_dqab=stride_dqab,
+        stride_dqah=stride_dqah,
+        stride_dqam=stride_dqam,
+        stride_dkb=stride_dkb,
+        stride_dkh=stride_dkh,
+        stride_dkn=stride_dkn,
+        stride_dks=stride_dks,
+        stride_dvb=stride_dvb,
+        stride_dvh=stride_dvh,
+        stride_dvn=stride_dvn,
+        stride_dvs=stride_dvs,
+        stride_dab=stride_dab,
+        stride_dah=stride_dah,
+        stride_ddb=stride_ddb,
+        stride_ddh=stride_ddh,
+        stride_dds=stride_dds,
+        IS_SPLIT_QO=IS_SPLIT_QO,
+        IS_GATED=True,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+    )
+    # Create block scheduler
+    block_sched = AttnBwdBlockScheduler.create(
+        config=config,
+        split_idx=grid_idx.split_idx,
+        num_splits=num_splits,
+        IS_CAUSAL=IS_CAUSAL,
+        IS_LOCAL=IS_LOCAL,
+        IS_SPLIT_QO=IS_SPLIT_QO,
+    )
+    # Create mask scheduler
+    mask_sched = AttnMaskScheduler.create(config, SWAP_AB=True)
+    # Create softmax scheduler
+    softmax_sched = SoftmaxScheduler.create(config)
+    # Early exit if no n_blocks to process
+    if grid_idx.n_block * TILE_N >= config.actual_seqlen_k:
+        return
+    # Create pointers
+    q_ptrs = ptrs_sched.make_q_ptrs(config)
+    k_ptrs = ptrs_sched.make_k_ptrs(config)
+    v_ptrs = ptrs_sched.make_v_ptrs(config)
+    a_ptrs = ptrs_sched.make_a_ptrs(config)
+    d_ptrs = ptrs_sched.make_d_ptrs(config)
+    do_ptrs = ptrs_sched.make_do_ptrs(config)
+    lse_ptrs = ptrs_sched.make_lse_ptrs(config)
+    dpsum_ptrs = ptrs_sched.make_dpsum_ptrs(config)
+    dq_accum_ptrs = ptrs_sched.make_dq_accum_ptrs(config)
+    dk_ptrs = ptrs_sched.make_dk_ptrs(config)
+    dv_ptrs = ptrs_sched.make_dv_ptrs(config)
+    da_ptrs = ptrs_sched.make_da_ptrs(config)
+    dd_ptrs = ptrs_sched.make_dd_ptrs(config)
+    # Initialize accumulators
+    gate_max = tl.full((), float("-inf"), dtype=tl.float32)
+    acc_dk = tl.zeros((TILE_N, TILE_K), dtype=tl.float32)
+    acc_dv = tl.zeros((TILE_N, TILE_K), dtype=tl.float32)
+    acc_dd = tl.zeros((TILE_N,), dtype=tl.float32)
+    # Load key tile
+    k_tile = ptrs_sched.load_k(config, k_ptrs)
+    # Load value tile
+    v_tile = ptrs_sched.load_v(config, v_ptrs)
+    # Load delta tile
+    d_tile = ptrs_sched.load_d(config, d_ptrs)
+    d_max = tl.max(d_tile)
+    d_min = tl.min(d_tile)
+    # Process m_blocks with causal masking
+    if IS_CAUSAL or IS_LOCAL:
+        for m_block in tl.range(
+            block_sched.m_block_min, block_sched.m_block_min_no_mask
+        ):
+            (
+                acc_dk,
+                acc_dv,
+                acc_dd,
+                gate_max,
+            ) = _bwd_inner_gated_kernel(
+                config=config,
+                ptrs_sched=ptrs_sched,
+                mask_sched=mask_sched,
+                softmax_sched=softmax_sched,
+                acc_dk=acc_dk,
+                acc_dv=acc_dv,
+                acc_dd=acc_dd,
+                k_tile=k_tile,
+                v_tile=v_tile,
+                d_tile=d_tile,
+                q_ptrs=q_ptrs,
+                a_ptrs=a_ptrs,
+                do_ptrs=do_ptrs,
+                dq_accum_ptrs=dq_accum_ptrs,
+                da_ptrs=da_ptrs,
+                lse_ptrs=lse_ptrs,
+                dpsum_ptrs=dpsum_ptrs,
+                d_max=d_max,
+                d_min=d_min,
+                gate_max=gate_max,
+                m_block=m_block,
+                IS_MASK=True,
+                MASK_CAUSAL=IS_CAUSAL,
+                MASK_LOCAL=True if IS_LOCAL else False,
+                MASK_SINK=False,
+            )
+    # Process m_blocks without masking
+    if not IS_LOCAL and block_sched.m_block_min_no_mask < block_sched.m_block_max:
+        for m_block in tl.range(
+            block_sched.m_block_min_no_mask, block_sched.m_block_max
+        ):
+            (
+                acc_dk,
+                acc_dv,
+                acc_dd,
+                gate_max,
+            ) = _bwd_inner_gated_kernel(
+                config=config,
+                ptrs_sched=ptrs_sched,
+                mask_sched=mask_sched,
+                softmax_sched=softmax_sched,
+                acc_dk=acc_dk,
+                acc_dv=acc_dv,
+                acc_dd=acc_dd,
+                k_tile=k_tile,
+                v_tile=v_tile,
+                d_tile=d_tile,
+                q_ptrs=q_ptrs,
+                a_ptrs=a_ptrs,
+                do_ptrs=do_ptrs,
+                dq_accum_ptrs=dq_accum_ptrs,
+                da_ptrs=da_ptrs,
+                lse_ptrs=lse_ptrs,
+                dpsum_ptrs=dpsum_ptrs,
+                d_max=d_max,
+                d_min=d_min,
+                gate_max=gate_max,
+                m_block=m_block,
+                IS_MASK=False,
+                MASK_CAUSAL=False,
+                MASK_LOCAL=False,
+                MASK_SINK=False,
+            )
+    if IS_LOCAL:
+        # Process m_blocks with local right masking
+        if block_sched.m_block_window_min < block_sched.m_block_window_min_no_mask:
+            for m_block in tl.range(
+                block_sched.m_block_window_min,
+                block_sched.m_block_window_min_no_mask,
+            ):
+                (
+                    acc_dk,
+                    acc_dv,
+                    acc_dd,
+                    gate_max,
+                ) = _bwd_inner_gated_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    acc_dk=acc_dk,
+                    acc_dv=acc_dv,
+                    acc_dd=acc_dd,
+                    k_tile=k_tile,
+                    v_tile=v_tile,
+                    d_tile=d_tile,
+                    q_ptrs=q_ptrs,
+                    a_ptrs=a_ptrs,
+                    do_ptrs=do_ptrs,
+                    dq_accum_ptrs=dq_accum_ptrs,
+                    da_ptrs=da_ptrs,
+                    lse_ptrs=lse_ptrs,
+                    dpsum_ptrs=dpsum_ptrs,
+                    d_max=d_max,
+                    d_min=d_min,
+                    gate_max=gate_max,
+                    m_block=m_block,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                )
+        # Process m_blocks without masking
+        if (
+            block_sched.m_block_window_min_no_mask
+            < block_sched.m_block_window_max_no_mask
+        ):
+            for m_block in tl.range(
+                block_sched.m_block_window_min_no_mask,
+                block_sched.m_block_window_max_no_mask,
+            ):
+                (
+                    acc_dk,
+                    acc_dv,
+                    acc_dd,
+                    gate_max,
+                ) = _bwd_inner_gated_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    acc_dk=acc_dk,
+                    acc_dv=acc_dv,
+                    acc_dd=acc_dd,
+                    k_tile=k_tile,
+                    v_tile=v_tile,
+                    d_tile=d_tile,
+                    q_ptrs=q_ptrs,
+                    a_ptrs=a_ptrs,
+                    do_ptrs=do_ptrs,
+                    dq_accum_ptrs=dq_accum_ptrs,
+                    da_ptrs=da_ptrs,
+                    lse_ptrs=lse_ptrs,
+                    dpsum_ptrs=dpsum_ptrs,
+                    d_max=d_max,
+                    d_min=d_min,
+                    gate_max=gate_max,
+                    m_block=m_block,
+                    IS_MASK=False,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=False,
+                    MASK_SINK=False,
+                )
+        # Process m_blocks with local left masking
+        if block_sched.m_block_window_max_no_mask < block_sched.m_block_window_max:
+            for m_block in tl.range(
+                block_sched.m_block_window_max_no_mask,
+                block_sched.m_block_window_max,
+            ):
+                (
+                    acc_dk,
+                    acc_dv,
+                    acc_dd,
+                    gate_max,
+                ) = _bwd_inner_gated_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    acc_dk=acc_dk,
+                    acc_dv=acc_dv,
+                    acc_dd=acc_dd,
+                    k_tile=k_tile,
+                    v_tile=v_tile,
+                    d_tile=d_tile,
+                    q_ptrs=q_ptrs,
+                    a_ptrs=a_ptrs,
+                    do_ptrs=do_ptrs,
+                    dq_accum_ptrs=dq_accum_ptrs,
+                    da_ptrs=da_ptrs,
+                    lse_ptrs=lse_ptrs,
+                    dpsum_ptrs=dpsum_ptrs,
+                    d_max=d_max,
+                    d_min=d_min,
+                    gate_max=gate_max,
+                    m_block=m_block,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                )
+        # Process m_blocks with local sink masking
+        if block_sched.m_block_sink_max > block_sched.m_block_sink_min:
+            for m_block in tl.range(
+                block_sched.m_block_sink_min,
+                block_sched.m_block_sink_max,
+            ):
+                (
+                    acc_dk,
+                    acc_dv,
+                    acc_dd,
+                    gate_max,
+                ) = _bwd_inner_gated_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    acc_dk=acc_dk,
+                    acc_dv=acc_dv,
+                    acc_dd=acc_dd,
+                    k_tile=k_tile,
+                    v_tile=v_tile,
+                    d_tile=d_tile,
+                    q_ptrs=q_ptrs,
+                    a_ptrs=a_ptrs,
+                    do_ptrs=do_ptrs,
+                    dq_accum_ptrs=dq_accum_ptrs,
+                    da_ptrs=da_ptrs,
+                    lse_ptrs=lse_ptrs,
+                    dpsum_ptrs=dpsum_ptrs,
+                    d_max=d_max,
+                    d_min=d_min,
+                    gate_max=gate_max,
+                    m_block=m_block,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=True,
+                )
+    # Store value gradients
+    ptrs_sched.store_dv(config, dv_ptrs, acc_dv)
+    # Scale delta gradients
+    acc_dd = acc_dd * softmax_scale
+    # Store delta gradients
+    ptrs_sched.store_dd(config, dd_ptrs, acc_dd)
+    # Scale key gradients
+    acc_dk = acc_dk * softmax_scale
+    # Store key gradients
+    ptrs_sched.store_dk(config, dk_ptrs, acc_dk)
+_bwd_gated_kernel = cache_utils.wrap_kernel(_bwd_gated_kernel)
+_bwd_gated_kernel_autotuned = None
+def _get_autotuned_kernel():
+    global _bwd_gated_kernel_autotuned
+    if _bwd_gated_kernel_autotuned is None:
+        jit_kernel = _bwd_gated_kernel._kernel
+        autotuned = autotuner.make_bwd_gated_autotuned_kernel(jit_kernel)
+        _bwd_gated_kernel_autotuned = autotuner.AutotunedKernel(autotuned)
+    return _bwd_gated_kernel_autotuned
+def _flash_gated_attn_backward(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    alpha: torch.Tensor,
+    delta: torch.Tensor,
+    out: torch.Tensor,
+    dout: torch.Tensor,
+    lse: torch.Tensor,
+    is_causal: bool = False,
+    softmax_scale: float = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    softmax_threshold: float = None,
+    gate_threshold: float = None,
+    is_logsigmoid_gate: bool = True,
+    is_adapt_gate: bool = True,
+    is_local: bool = False,
+    is_quant: bool = False,
+    is_split_qo: bool = False,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    batch_size, seqlen_q, num_heads_q, head_dim = query.shape
+    _, seqlen_k, num_heads_kv, _ = key.shape
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    softmax_threshold = (
+        softmax_threshold if softmax_threshold is not None else 1 / seqlen_k
+    )
+    gate_threshold = (
+        gate_threshold if gate_threshold is not None else head_dim / seqlen_k
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_bwd_inputs(
+            query,
+            key,
+            value,
+            out,
+            dout,
+            lse,
+            alpha=alpha,
+            delta=delta,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            cu_seqlens_q=None,
+            cu_seqlens_k=None,
+            seqused_q=None,
+            seqused_k=None,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name="bwd_gated",
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_causal=is_causal,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _bwd_gated_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        # Placeholder for pre-launch computations
+        TILE_M = TILE_N = 64
+        num_warps = num_stages = num_ctas = None
+    num_splits = (
+        utils.num_splits_heuristic(
+            seqlen_q=seqlen_k,
+            seqlen_k=seqlen_q,
+            num_SMs=num_SMs,
+            TILE_M=TILE_N,
+            TILE_N=TILE_M,
+            max_split_blocks=utils.max_split_blocks_from_window_sizes(
+                window_sizes, TILE_M
+            )
+            if is_local
+            else None,
+        )
+        if is_split_qo
+        else 1
+    )
+    seqlen_q_rounded = int(math.ceil(seqlen_q / 128) * 128)
+    head_dim_rounded = int(math.ceil(head_dim / 32) * 32)
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    dq = torch.empty_like(query, dtype=query_scale.dtype)
+    dk = torch.empty_like(key, dtype=key_scale.dtype)
+    dv = torch.empty_like(value, dtype=value_scale.dtype)
+    da = torch.empty_like(alpha)
+    dd = torch.empty_like(delta)
+    lse_log2 = torch.empty(
+        (batch_size, num_heads_q, seqlen_q_rounded),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dpsum = torch.empty(
+        (batch_size, num_heads_q, seqlen_q_rounded),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dq_accum = torch.empty(
+        (batch_size, num_heads_q, seqlen_q_rounded * head_dim_rounded),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dk_accum = torch.zeros(
+        (num_splits, batch_size, seqlen_k, num_heads_kv, head_dim)
+        if is_split_qo and num_splits > 1
+        else (batch_size, seqlen_k, num_heads_kv, head_dim),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dv_accum = torch.zeros(
+        (num_splits, batch_size, seqlen_k, num_heads_kv, head_dim)
+        if is_split_qo and num_splits > 1
+        else (batch_size, seqlen_k, num_heads_kv, head_dim),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    da_accum = torch.zeros(
+        (batch_size, num_heads_q, seqlen_q_rounded),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dd_accum = torch.zeros(
+        (num_splits, batch_size, num_heads_kv, seqlen_k)
+        if is_split_qo and num_splits > 1
+        else (batch_size, num_heads_kv, seqlen_k),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    flash_bwd_preprocess._flash_attn_bwd_preprocess(
+        out=out,
+        dout=dout,
+        dpsum=dpsum,
+        lse=lse,
+        lse_log2=lse_log2,
+        dq_accum=dq_accum,
+        head_dim_rounded=head_dim_rounded,
+        tile_m=TILE_M,
+        tile_k=TILE_K,
+    )
+    grid = launch_grid.get_bwd_grid(
+        seqlen_k=seqlen_k,
+        num_heads_q=num_heads_q,
+        batch_size=batch_size,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        alpha,
+        delta,
+        dout,
+        lse_log2,
+        dpsum,
+        dq_accum,
+        dk_accum,
+        dv_accum,
+        da_accum,
+        dd_accum,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        softmax_threshold,
+        gate_threshold,
+        query.stride(0),
+        query.stride(-2),
+        query.stride(-3),
+        key.stride(0),
+        key.stride(-2),
+        key.stride(-3),
+        value.stride(0),
+        value.stride(-2),
+        value.stride(-3),
+        alpha.stride(0),
+        alpha.stride(-2),
+        delta.stride(0),
+        delta.stride(-2),
+        dout.stride(0),
+        dout.stride(-2),
+        dout.stride(-3),
+        lse_log2.stride(0),
+        lse_log2.stride(-2),
+        dpsum.stride(0),
+        dpsum.stride(-2),
+        dq_accum.stride(0),
+        dq_accum.stride(1),
+        head_dim_rounded,
+        dk_accum.stride(-4) if is_split_qo and num_splits > 1 else dk_accum.stride(0),
+        dk_accum.stride(-2),
+        dk_accum.stride(-3),
+        dk_accum.stride(0) if is_split_qo and num_splits > 1 else 0,
+        dv_accum.stride(-4) if is_split_qo and num_splits > 1 else dv_accum.stride(0),
+        dv_accum.stride(-2),
+        dv_accum.stride(-3),
+        dv_accum.stride(0) if is_split_qo and num_splits > 1 else 0,
+        da_accum.stride(0),
+        da_accum.stride(-2),
+        dd_accum.stride(-3) if is_split_qo and num_splits > 1 else dd_accum.stride(0),
+        dd_accum.stride(-2),
+        dd_accum.stride(0) if is_split_qo and num_splits > 1 else 0,
+        window_sizes.stride(0),
+        None,
+        None,
+        None,
+        None,
+        num_splits,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=seqlen_q // 1024,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD=qhead_per_kvhead,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_CAUSAL=is_causal,
+        IS_LOCAL=is_local,
+        IS_SPLIT_QO=is_split_qo and num_splits > 1,
+        IS_LOGSIGMOID_GATE=is_logsigmoid_gate,
+        IS_ADAPT_GATE=is_adapt_gate,
+        HAS_CU_SEQLENS_Q=False,
+        HAS_CU_SEQLENS_K=False,
+        HAS_SEQUSED_Q=False,
+        HAS_SEQUSED_K=False,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name="bwd_gated",
+                seqlen_q=seqlen_q,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_causal=is_causal,
+                is_quant=is_quant,
+            )
+    flash_bwd_postprocess._flash_attn_bwd_postprocess(
+        dq_accum=dq_accum,
+        dq=dq,
+        da_accum=da_accum,
+        da=da,
+        scale=softmax_scale,
+        head_dim_rounded=head_dim_rounded,
+        tile_m=TILE_M,
+        tile_k=TILE_K,
+    )
+    if is_split_qo and num_splits > 1:
+        dk_accum = dk_accum.sum(dim=0)
+        dv_accum = dv_accum.sum(dim=0)
+        dd_accum = dd_accum.sum(dim=0)
+    dk.copy_(dk_accum)
+    dv.copy_(dv_accum)
+    dd.copy_(dd_accum)
+    return dq, dk, dv, da, dd
+def _flash_gated_attn_varlen_backward(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    alpha: torch.Tensor,
+    delta: torch.Tensor,
+    out: torch.Tensor,
+    dout: torch.Tensor,
+    lse: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    is_causal: bool = False,
+    softmax_scale: float = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    softmax_threshold: float = None,
+    gate_threshold: float = None,
+    is_logsigmoid_gate: bool = True,
+    is_adapt_gate: bool = True,
+    is_local: bool = False,
+    is_quant: bool = False,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    is_split_qo: bool = False,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    total_q, num_heads_q, head_dim = query.shape
+    total_k, num_heads_kv, _ = key.shape
+    batch_size = cu_seqlens_q.shape[0] - 1
+    seqlen_q = max_seqlen_q
+    seqlen_k = max_seqlen_k
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    softmax_threshold = (
+        softmax_threshold if softmax_threshold is not None else 1 / seqlen_k
+    )
+    gate_threshold = (
+        gate_threshold if gate_threshold is not None else head_dim / seqlen_k
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_bwd_inputs(
+            query,
+            key,
+            value,
+            out,
+            dout,
+            lse,
+            alpha=alpha,
+            delta=delta,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name="bwd_gated",
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_causal=is_causal,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _bwd_gated_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        # Placeholder for pre-launch computations
+        TILE_M = TILE_N = 64
+        num_warps = num_stages = num_ctas = None
+    num_splits = (
+        utils.num_splits_heuristic(
+            seqlen_q=seqlen_k,
+            seqlen_k=seqlen_q,
+            num_SMs=num_SMs,
+            TILE_M=TILE_N,
+            TILE_N=TILE_M,
+            max_split_blocks=utils.max_split_blocks_from_window_sizes(
+                window_sizes, TILE_M
+            )
+            if is_local
+            else None,
+        )
+        if is_split_qo
+        else 1
+    )
+    total_q_rounded_padded = int(math.ceil((total_q + batch_size * 128) / 128) * 128)
+    head_dim_rounded = int(math.ceil(head_dim / 32) * 32)
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    dq = torch.empty_like(query, dtype=query_scale.dtype)
+    dk = torch.empty_like(key, dtype=key_scale.dtype)
+    dv = torch.empty_like(value, dtype=value_scale.dtype)
+    da = torch.empty_like(alpha)
+    dd = torch.empty_like(delta)
+    lse_log2 = torch.empty(
+        (num_heads_q, total_q_rounded_padded),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dpsum = torch.empty(
+        (num_heads_q, total_q_rounded_padded),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dq_accum = torch.empty(
+        (num_heads_q, total_q_rounded_padded * head_dim_rounded),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dk_accum = torch.zeros(
+        (num_splits, total_k, num_heads_kv, head_dim)
+        if is_split_qo and num_splits > 1
+        else (total_k, num_heads_kv, head_dim),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dv_accum = torch.zeros(
+        (num_splits, total_k, num_heads_kv, head_dim)
+        if is_split_qo and num_splits > 1
+        else (total_k, num_heads_kv, head_dim),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    da_accum = torch.zeros(
+        (num_heads_q, total_q_rounded_padded),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dd_accum = torch.zeros(
+        (num_splits, num_heads_kv, total_k)
+        if is_split_qo and num_splits > 1
+        else (num_heads_kv, total_k),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    flash_bwd_preprocess._flash_attn_bwd_preprocess(
+        out=out,
+        dout=dout,
+        dpsum=dpsum,
+        lse=lse,
+        lse_log2=lse_log2,
+        dq_accum=dq_accum,
+        head_dim_rounded=head_dim_rounded,
+        cu_seqlens_q=cu_seqlens_q,
+        seqused_q=seqused_q,
+        max_seqlen_q=max_seqlen_q,
+        tile_m=TILE_M,
+        tile_k=TILE_K,
+    )
+    grid = launch_grid.get_bwd_grid(
+        seqlen_k=seqlen_k,
+        num_heads_q=num_heads_q,
+        batch_size=batch_size,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        alpha,
+        delta,
+        dout,
+        lse_log2,
+        dpsum,
+        dq_accum,
+        dk_accum,
+        dv_accum,
+        da_accum,
+        dd_accum,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        softmax_threshold,
+        gate_threshold,
+        0,
+        query.stride(-2),
+        query.stride(0),
+        0,
+        key.stride(-2),
+        key.stride(0),
+        0,
+        value.stride(-2),
+        value.stride(0),
+        0,
+        alpha.stride(-2),
+        0,
+        delta.stride(-2),
+        0,
+        dout.stride(-2),
+        dout.stride(0),
+        0,
+        lse_log2.stride(0),
+        0,
+        dpsum.stride(0),
+        0,
+        dq_accum.stride(0),
+        head_dim_rounded,
+        0,
+        dk_accum.stride(-2),
+        dk_accum.stride(-3),
+        dk_accum.stride(0) if is_split_qo and num_splits > 1 else 0,
+        0,
+        dv_accum.stride(-2),
+        dv_accum.stride(-3),
+        dv_accum.stride(0) if is_split_qo and num_splits > 1 else 0,
+        0,
+        da_accum.stride(0),
+        0,
+        dd_accum.stride(-2),
+        dd_accum.stride(0) if is_split_qo and num_splits > 1 else 0,
+        window_sizes.stride(0),
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_q,
+        seqused_k,
+        num_splits,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=seqlen_q // 1024,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD=qhead_per_kvhead,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_CAUSAL=is_causal,
+        IS_LOCAL=is_local,
+        IS_SPLIT_QO=is_split_qo and num_splits > 1,
+        IS_LOGSIGMOID_GATE=is_logsigmoid_gate,
+        IS_ADAPT_GATE=is_adapt_gate,
+        HAS_CU_SEQLENS_Q=True,
+        HAS_CU_SEQLENS_K=True,
+        HAS_SEQUSED_Q=seqused_q is not None,
+        HAS_SEQUSED_K=seqused_k is not None,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name="bwd_gated",
+                seqlen_q=seqlen_q,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_causal=is_causal,
+                is_quant=is_quant,
+            )
+    flash_bwd_postprocess._flash_attn_bwd_postprocess(
+        dq_accum=dq_accum,
+        dq=dq,
+        da_accum=da_accum,
+        da=da,
+        scale=softmax_scale,
+        head_dim_rounded=head_dim_rounded,
+        cu_seqlens_q=cu_seqlens_q,
+        seqused_q=seqused_q,
+        max_seqlen_q=max_seqlen_q,
+        tile_m=TILE_M,
+        tile_k=TILE_K,
+    )
+    if is_split_qo and num_splits > 1:
+        dk_accum = dk_accum.sum(dim=0)
+        dv_accum = dv_accum.sum(dim=0)
+        dd_accum = dd_accum.sum(dim=0)
+    dk.copy_(dk_accum)
+    dv.copy_(dv_accum)
+    dd.copy_(dd_accum)
+    return dq, dk, dv, da, dd

build/torch-cuda/flash_gated_dec.py ADDED Viewed

	@@ -0,0 +1,1218 @@

+from typing import Tuple, Optional
+import torch
+import triton
+import triton.language as tl
+from . import (
+    assert_inputs,
+    utils,
+    cache_utils,
+    launch_template,
+    launch_grid,
+    flash_dec_combine,
+    kernel_repr,
+    autotuner,
+)
+from .scheduler import (
+    AttnDecGridIndex,
+    AttnDecConfig,
+    AttnDecBlockScheduler,
+    AttnDecPointerScheduler,
+    AttnMaskScheduler,
+    SoftmaxScheduler,
+)
+@triton.jit
+def _dec_inner_gated_kernel(
+    config: AttnDecConfig,
+    ptrs_sched: AttnDecPointerScheduler,
+    mask_sched: AttnMaskScheduler,
+    softmax_sched: SoftmaxScheduler,
+    skip_gate_curr,
+    acc_s,
+    acc_o,
+    q_tile,
+    a_tile,
+    k_ptrs,
+    v_ptrs,
+    d_ptrs,
+    a_max,
+    a_min,
+    gate_max,
+    row_max,
+    row_sum,
+    n_block,
+    n_block_min,
+    IS_MASK: tl.constexpr,
+    MASK_LOCAL: tl.constexpr,
+    MASK_SINK: tl.constexpr,
+    CHECK_INF: tl.constexpr,
+):
+    # Load next delta tile
+    d_tile = ptrs_sched.load_d(config, d_ptrs, n_block - 1)
+    d_max = tl.max(d_tile)
+    d_min = tl.min(d_tile)
+    skip_gate_next = True
+    if not skip_gate_curr:
+        if n_block > n_block_min:
+            # Check if any gates are active for next tile
+            gate_max, skip_gate_next = softmax_sched.online_gate(
+                a_max=a_max,
+                a_min=a_min,
+                d_max=d_max,
+                d_min=d_min,
+                gate_max=gate_max,
+                gate_threshold_log2=config.gate_threshold_log2,
+            )
+        if IS_MASK:
+            # Apply mask to attention scores
+            acc_s = mask_sched.apply_mask(
+                acc_s=acc_s,
+                iter_block=n_block,
+                MASK_LOCAL=MASK_LOCAL,
+                MASK_SINK=MASK_SINK,
+            )
+        # Apply online softmax
+        p, row_max, row_sum, row_scale, skip_softmax = (
+            softmax_sched.online_sparse_softmax(
+                acc_s=acc_s,
+                row_max=row_max,
+                row_sum=row_sum,
+                softmax_threshold_log2=config.softmax_threshold_log2,
+                CHECK_INF=CHECK_INF,
+            )
+        )
+        if not skip_softmax:
+            # Load value tile
+            v_tile = ptrs_sched.load_v(config, v_ptrs, n_block)
+            # Rescale output accumulator
+            acc_o = softmax_sched.rescale_o(
+                acc_o=acc_o,
+                row_scale=row_scale,
+            )
+            # Update output accumulator
+            acc_o += tl.dot(p.to(v_tile.dtype), v_tile)
+    else:
+        if n_block > n_block_min:
+            # Check if any gates are active for next tile
+            gate_max, skip_gate_next = softmax_sched.online_gate(
+                a_max=a_max,
+                a_min=a_min,
+                d_max=d_max,
+                d_min=d_min,
+                gate_max=gate_max,
+                gate_threshold_log2=config.gate_threshold_log2,
+            )
+    if not skip_gate_next:
+        # Compute attention gates for next tile
+        acc_s = a_tile[:, None] * d_tile[None, :]
+        if config.IS_LOGSIGMOID_GATE:
+            acc_s = softmax_sched.log_sigmoid(
+                acc_s=acc_s,
+            )
+        # Load next key tile
+        k_tile = ptrs_sched.load_k(config, k_ptrs, n_block - 1)
+        # Compute attention scores for next tile
+        acc_s += tl.dot(q_tile, k_tile.T)
+    return (
+        skip_gate_next,
+        acc_s,
+        acc_o,
+        gate_max,
+        row_max,
+        row_sum,
+    )
+@triton.jit(repr=kernel_repr.dec_gated_repr)
+def _dec_gated_kernel(
+    Q,
+    K,
+    V,
+    A,
+    D,
+    Out,
+    Lse,
+    softmax_scale,
+    query_scale,
+    key_scale,
+    value_scale,
+    softmax_threshold,
+    gate_threshold,
+    window_sizes,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_ab,
+    stride_ah,
+    stride_db,
+    stride_dh,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_os,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_ls,
+    stride_wh,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    seqused_q,
+    seqused_k,
+    num_splits,
+    seqlen_q,
+    seqlen_k,
+    head_dim,
+    SEQLEN_Q_CACHE: tl.constexpr,
+    SEQLEN_K_CACHE: tl.constexpr,
+    QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+    TILE_K: tl.constexpr,
+    IS_LOCAL: tl.constexpr,
+    HAS_CU_SEQLENS_Q: tl.constexpr,
+    HAS_CU_SEQLENS_K: tl.constexpr,
+    HAS_SEQUSED_Q: tl.constexpr,
+    HAS_SEQUSED_K: tl.constexpr,
+    IS_LOGSIGMOID_GATE: tl.constexpr,
+):
+    # Create grid index
+    grid_idx = AttnDecGridIndex.create(
+        num_splits=num_splits,
+    )
+    # Load window sizes
+    (
+        window_size_sink,
+        window_size_left,
+        window_size_right,
+        window_size_dist,
+    ) = grid_idx.load_window_sizes(
+        window_sizes=window_sizes,
+        stride_wh=stride_wh,
+        IS_LOCAL=IS_LOCAL,
+    )
+    # Create config
+    config = AttnDecConfig.create(
+        softmax_scale=softmax_scale,
+        softmax_threshold=softmax_threshold,
+        gate_threshold=gate_threshold,
+        query_scale=query_scale,
+        key_scale=key_scale,
+        value_scale=value_scale,
+        batch_idx=grid_idx.batch_idx,
+        window_size_sink=window_size_sink,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        window_size_dist=window_size_dist,
+        head_dim=head_dim,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        QHEAD_PER_KVHEAD_PACKGQA=QHEAD_PER_KVHEAD_PACKGQA,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_LOGSIGMOID_GATE=IS_LOGSIGMOID_GATE,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+        HAS_SEQUSED_Q=HAS_SEQUSED_Q,
+        HAS_SEQUSED_K=HAS_SEQUSED_K,
+    )
+    # Create pointer scheduler
+    ptrs_sched = AttnDecPointerScheduler.create(
+        config=config,
+        Q=Q,
+        K=K,
+        V=V,
+        A=A,
+        D=D,
+        Out=Out,
+        Lse=Lse,
+        batch_idx=grid_idx.batch_idx,
+        head_idx=grid_idx.head_idx,
+        head_kv_idx=grid_idx.head_kv_idx,
+        split_idx=grid_idx.split_idx,
+        stride_qb=stride_qb,
+        stride_qh=stride_qh,
+        stride_qm=stride_qm,
+        stride_kb=stride_kb,
+        stride_kh=stride_kh,
+        stride_kn=stride_kn,
+        stride_vb=stride_vb,
+        stride_vh=stride_vh,
+        stride_vn=stride_vn,
+        stride_ab=stride_ab,
+        stride_ah=stride_ah,
+        stride_db=stride_db,
+        stride_dh=stride_dh,
+        stride_ob=stride_ob,
+        stride_oh=stride_oh,
+        stride_om=stride_om,
+        stride_os=stride_os,
+        stride_lb=stride_lb,
+        stride_lh=stride_lh,
+        stride_lm=stride_lm,
+        stride_ls=stride_ls,
+        IS_GATED=True,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+    )
+    # Create block scheduler
+    block_sched = AttnDecBlockScheduler.create(
+        config=config,
+        split_idx=grid_idx.split_idx,
+        num_splits=num_splits,
+        IS_LOCAL=IS_LOCAL,
+    )
+    # Create mask scheduler
+    mask_sched = AttnMaskScheduler.create(config)
+    # Create softmax scheduler
+    softmax_sched = SoftmaxScheduler.create(config)
+    # Create pointers
+    out_ptrs = ptrs_sched.make_out_ptrs(config)
+    lse_ptrs = ptrs_sched.make_lse_ptrs(config)
+    # Early exit if no n_blocks to process
+    if block_sched.is_empty():
+        ptrs_sched.store_empty(config, out_ptrs, lse_ptrs, Out)
+        return
+    q_ptrs = ptrs_sched.make_q_ptrs(config)
+    k_ptrs = ptrs_sched.make_k_ptrs(config)
+    v_ptrs = ptrs_sched.make_v_ptrs(config)
+    a_ptrs = ptrs_sched.make_a_ptrs(config)
+    d_ptrs = ptrs_sched.make_d_ptrs(config)
+    # Initialize accumulators
+    gate_max = tl.full((), float("-inf"), dtype=tl.float32)
+    row_max = tl.full((TILE_M,), float("-inf"), dtype=tl.float32)
+    row_sum = tl.zeros((TILE_M,), dtype=tl.float32)
+    acc_o = tl.zeros((TILE_M, TILE_K), dtype=tl.float32)
+    # Load query tile
+    q_tile = ptrs_sched.load_q(config, q_ptrs)
+    # Load alpha tile
+    a_tile = ptrs_sched.load_a(config, a_ptrs)
+    a_max = tl.max(a_tile)
+    a_min = tl.min(a_tile)
+    # Load delta tile
+    d_tile = ptrs_sched.load_d(config, d_ptrs, block_sched.n_block_max - 1)
+    d_max = tl.max(d_tile)
+    d_min = tl.min(d_tile)
+    # Check if any gates are active for current tile
+    gate_max, skip_gate_curr = softmax_sched.online_gate(
+        a_max=a_max,
+        a_min=a_min,
+        d_max=d_max,
+        d_min=d_min,
+        gate_max=gate_max,
+        gate_threshold_log2=config.gate_threshold_log2,
+    )
+    # Initialize skip_gate_curr to False for the first iteration
+    skip_gate_curr = False
+    # Compute attention gates for first tile
+    acc_s = a_tile[:, None] * d_tile[None, :]
+    if IS_LOGSIGMOID_GATE:
+        acc_s = softmax_sched.log_sigmoid(
+            acc_s=acc_s,
+        )
+    # Load key tile
+    k_tile = ptrs_sched.load_k(config, k_ptrs, block_sched.n_block_max - 1)
+    # Compute attention scores for first tile
+    acc_s += tl.dot(q_tile, k_tile.T)
+    # Process n_blocks with seqlen masking
+    for n_block in tl.range(
+        block_sched.n_block_max - 1, block_sched.n_block_max_no_mask - 1, -1
+    ):
+        (
+            skip_gate_curr,
+            acc_s,
+            acc_o,
+            gate_max,
+            row_max,
+            row_sum,
+        ) = _dec_inner_gated_kernel(
+            config=config,
+            ptrs_sched=ptrs_sched,
+            mask_sched=mask_sched,
+            softmax_sched=softmax_sched,
+            skip_gate_curr=skip_gate_curr,
+            acc_s=acc_s,
+            acc_o=acc_o,
+            q_tile=q_tile,
+            a_tile=a_tile,
+            k_ptrs=k_ptrs,
+            v_ptrs=v_ptrs,
+            d_ptrs=d_ptrs,
+            a_max=a_max,
+            a_min=a_min,
+            gate_max=gate_max,
+            row_max=row_max,
+            row_sum=row_sum,
+            n_block=n_block,
+            n_block_min=block_sched.n_block_max_no_mask,
+            IS_MASK=True,
+            MASK_LOCAL=True if IS_LOCAL else False,
+            MASK_SINK=False,
+            CHECK_INF=True,
+        )
+    # Process n_blocks without masking
+    if not IS_LOCAL and block_sched.n_block_max_no_mask > block_sched.n_block_min:
+        # Load key tile
+        k_tile = ptrs_sched.load_k(config, k_ptrs, block_sched.n_block_max_no_mask - 1)
+        # Load delta tile
+        d_tile = ptrs_sched.load_d(config, d_ptrs, block_sched.n_block_max_no_mask - 1)
+        d_max = tl.max(d_tile)
+        d_min = tl.min(d_tile)
+        # Check if any gates are active for current tile
+        gate_max, skip_gate_curr = softmax_sched.online_gate(
+            a_max=a_max,
+            a_min=a_min,
+            d_max=d_max,
+            d_min=d_min,
+            gate_max=gate_max,
+            gate_threshold_log2=config.gate_threshold_log2,
+        )
+        # Compute attention gates
+        acc_s = a_tile[:, None] * d_tile[None, :]
+        if IS_LOGSIGMOID_GATE:
+            acc_s = softmax_sched.log_sigmoid(
+                acc_s=acc_s,
+            )
+        # Compute attention scores
+        acc_s += tl.dot(q_tile, k_tile.T)
+        for n_block in tl.range(
+            block_sched.n_block_max_no_mask - 1, block_sched.n_block_min - 1, -1
+        ):
+            (
+                skip_gate_curr,
+                acc_s,
+                acc_o,
+                gate_max,
+                row_max,
+                row_sum,
+            ) = _dec_inner_gated_kernel(
+                config=config,
+                ptrs_sched=ptrs_sched,
+                mask_sched=mask_sched,
+                softmax_sched=softmax_sched,
+                skip_gate_curr=skip_gate_curr,
+                acc_s=acc_s,
+                acc_o=acc_o,
+                q_tile=q_tile,
+                a_tile=a_tile,
+                k_ptrs=k_ptrs,
+                v_ptrs=v_ptrs,
+                d_ptrs=d_ptrs,
+                a_max=a_max,
+                a_min=a_min,
+                gate_max=gate_max,
+                row_max=row_max,
+                row_sum=row_sum,
+                n_block=n_block,
+                n_block_min=block_sched.n_block_min,
+                IS_MASK=False,
+                MASK_LOCAL=False,
+                MASK_SINK=False,
+                CHECK_INF=False,
+            )
+    if IS_LOCAL:
+        # Process n_blocks with local right masking
+        if block_sched.n_block_window_max > block_sched.n_block_window_max_no_mask:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_max - 1
+            )
+            # Load delta tile
+            d_tile = ptrs_sched.load_d(
+                config, d_ptrs, block_sched.n_block_window_max - 1
+            )
+            d_max = tl.max(d_tile)
+            d_min = tl.min(d_tile)
+            # Check if any gates are active for current tile
+            gate_max, skip_gate_curr = softmax_sched.online_gate(
+                a_max=a_max,
+                a_min=a_min,
+                d_max=d_max,
+                d_min=d_min,
+                gate_max=gate_max,
+                gate_threshold_log2=config.gate_threshold_log2,
+            )
+            # Compute attention gates
+            acc_s = a_tile[:, None] * d_tile[None, :]
+            if IS_LOGSIGMOID_GATE:
+                acc_s = softmax_sched.log_sigmoid(
+                    acc_s=acc_s,
+                )
+            # Compute attention scores
+            acc_s += tl.dot(q_tile, k_tile.T)
+            for n_block in tl.range(
+                block_sched.n_block_window_max - 1,
+                block_sched.n_block_window_max_no_mask - 1,
+                -1,
+            ):
+                (
+                    skip_gate_curr,
+                    acc_s,
+                    acc_o,
+                    gate_max,
+                    row_max,
+                    row_sum,
+                ) = _dec_inner_gated_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    skip_gate_curr=skip_gate_curr,
+                    acc_s=acc_s,
+                    acc_o=acc_o,
+                    q_tile=q_tile,
+                    a_tile=a_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    d_ptrs=d_ptrs,
+                    a_max=a_max,
+                    a_min=a_min,
+                    gate_max=gate_max,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_max_no_mask,
+                    IS_MASK=True,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                    CHECK_INF=True,
+                )
+        # Process n_blocks without masking
+        if (
+            block_sched.n_block_window_max_no_mask
+            > block_sched.n_block_window_min_no_mask
+        ):
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_max_no_mask - 1
+            )
+            # Load delta tile
+            d_tile = ptrs_sched.load_d(
+                config, d_ptrs, block_sched.n_block_window_max_no_mask - 1
+            )
+            d_max = tl.max(d_tile)
+            d_min = tl.min(d_tile)
+            # Check if any gates are active for current tile
+            gate_max, skip_gate_curr = softmax_sched.online_gate(
+                a_max=a_max,
+                a_min=a_min,
+                d_max=d_max,
+                d_min=d_min,
+                gate_max=gate_max,
+                gate_threshold_log2=config.gate_threshold_log2,
+            )
+            # Compute attention gates
+            acc_s = a_tile[:, None] * d_tile[None, :]
+            if IS_LOGSIGMOID_GATE:
+                acc_s = softmax_sched.log_sigmoid(
+                    acc_s=acc_s,
+                )
+            # Compute attention scores
+            acc_s += tl.dot(q_tile, k_tile.T)
+            for n_block in tl.range(
+                block_sched.n_block_window_max_no_mask - 1,
+                block_sched.n_block_window_min_no_mask - 1,
+                -1,
+            ):
+                (
+                    skip_gate_curr,
+                    acc_s,
+                    acc_o,
+                    gate_max,
+                    row_max,
+                    row_sum,
+                ) = _dec_inner_gated_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    skip_gate_curr=skip_gate_curr,
+                    acc_s=acc_s,
+                    acc_o=acc_o,
+                    q_tile=q_tile,
+                    a_tile=a_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    d_ptrs=d_ptrs,
+                    a_max=a_max,
+                    a_min=a_min,
+                    gate_max=gate_max,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_min_no_mask,
+                    IS_MASK=False,
+                    MASK_LOCAL=False,
+                    MASK_SINK=False,
+                    CHECK_INF=False,
+                )
+        # Process n_blocks with local left masking
+        if block_sched.n_block_window_min_no_mask > block_sched.n_block_window_min:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_min_no_mask - 1
+            )
+            # Load delta tile
+            d_tile = ptrs_sched.load_d(
+                config, d_ptrs, block_sched.n_block_window_min_no_mask - 1
+            )
+            d_max = tl.max(d_tile)
+            d_min = tl.min(d_tile)
+            # Check if any gates are active for current tile
+            gate_max, skip_gate_curr = softmax_sched.online_gate(
+                a_max=a_max,
+                a_min=a_min,
+                d_max=d_max,
+                d_min=d_min,
+                gate_max=gate_max,
+                gate_threshold_log2=config.gate_threshold_log2,
+            )
+            # Compute attention gates
+            acc_s = a_tile[:, None] * d_tile[None, :]
+            if IS_LOGSIGMOID_GATE:
+                acc_s = softmax_sched.log_sigmoid(
+                    acc_s=acc_s,
+                )
+            # Compute attention scores
+            acc_s += tl.dot(q_tile, k_tile.T)
+            for n_block in tl.range(
+                block_sched.n_block_window_min_no_mask - 1,
+                block_sched.n_block_window_min - 1,
+                -1,
+            ):
+                (
+                    skip_gate_curr,
+                    acc_s,
+                    acc_o,
+                    gate_max,
+                    row_max,
+                    row_sum,
+                ) = _dec_inner_gated_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    skip_gate_curr=skip_gate_curr,
+                    acc_s=acc_s,
+                    acc_o=acc_o,
+                    q_tile=q_tile,
+                    a_tile=a_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    d_ptrs=d_ptrs,
+                    a_max=a_max,
+                    a_min=a_min,
+                    gate_max=gate_max,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_min,
+                    IS_MASK=True,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                    CHECK_INF=True,
+                )
+        # Process n_blocks with local sink masking
+        if block_sched.n_block_sink_max > block_sched.n_block_sink_min:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(config, k_ptrs, block_sched.n_block_sink_max - 1)
+            # Load delta tile
+            d_tile = ptrs_sched.load_d(config, d_ptrs, block_sched.n_block_sink_max - 1)
+            d_max = tl.max(d_tile)
+            d_min = tl.min(d_tile)
+            # Check if any gates are active for current tile
+            gate_max, skip_gate_curr = softmax_sched.online_gate(
+                a_max=a_max,
+                a_min=a_min,
+                d_max=d_max,
+                d_min=d_min,
+                gate_max=gate_max,
+                gate_threshold_log2=config.gate_threshold_log2,
+            )
+            # Compute attention gates
+            acc_s = a_tile[:, None] * d_tile[None, :]
+            if IS_LOGSIGMOID_GATE:
+                acc_s = softmax_sched.log_sigmoid(
+                    acc_s=acc_s,
+                )
+            # Compute attention scores
+            acc_s += tl.dot(q_tile, k_tile.T)
+            for n_block in tl.range(
+                block_sched.n_block_sink_max - 1,
+                block_sched.n_block_sink_min - 1,
+                -1,
+            ):
+                (
+                    skip_gate_curr,
+                    acc_s,
+                    acc_o,
+                    gate_max,
+                    row_max,
+                    row_sum,
+                ) = _dec_inner_gated_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    skip_gate_curr=skip_gate_curr,
+                    acc_s=acc_s,
+                    acc_o=acc_o,
+                    q_tile=q_tile,
+                    a_tile=a_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    d_ptrs=d_ptrs,
+                    a_max=a_max,
+                    a_min=a_min,
+                    gate_max=gate_max,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_sink_min,
+                    IS_MASK=True,
+                    MASK_LOCAL=True,
+                    MASK_SINK=True,
+                    CHECK_INF=True,
+                )
+    # Finalize softmax
+    row_scale, lse_tile = softmax_sched.finalize(
+        row_max=row_max,
+        row_sum=row_sum,
+        IS_LOG2=True,
+    )
+    # Store LSE
+    ptrs_sched.store_lse(config, lse_ptrs, lse_tile)
+    # Finalize rescale
+    acc_o = softmax_sched.rescale_o(
+        acc_o=acc_o,
+        row_scale=row_scale,
+    )
+    # Store output
+    ptrs_sched.store_out(config, out_ptrs, acc_o)
+_dec_gated_kernel = cache_utils.wrap_kernel(_dec_gated_kernel)
+_dec_gated_kernel_autotuned = None
+def _get_autotuned_kernel():
+    global _dec_gated_kernel_autotuned
+    if _dec_gated_kernel_autotuned is None:
+        jit_kernel = _dec_gated_kernel._kernel
+        autotuned = autotuner.make_dec_gated_autotuned_kernel(jit_kernel)
+        _dec_gated_kernel_autotuned = autotuner.AutotunedKernel(autotuned)
+    return _dec_gated_kernel_autotuned
+def _flash_gated_attn_decode(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    alpha: torch.Tensor,
+    delta: torch.Tensor,
+    softmax_scale: float = None,
+    softmax_threshold: float = None,
+    gate_threshold: float = None,
+    is_logsigmoid_gate: bool = True,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    batch_size, num_heads_q, head_dim = query.shape
+    _, seqlen_k, num_heads_kv, _ = key.shape
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    softmax_threshold = (
+        softmax_threshold if softmax_threshold is not None else 1 / seqlen_k
+    )
+    gate_threshold = (
+        gate_threshold if gate_threshold is not None else head_dim / seqlen_k
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_dec_inputs(
+            query,
+            key,
+            value,
+            alpha=alpha,
+            delta=delta,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            cu_seqlens_k=None,
+            seqused_k=None,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name="dec_gated",
+        seqlen_q=1,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _dec_gated_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        TILE_M = max(triton.next_power_of_2(qhead_per_kvhead), 16)
+        TILE_N = 128
+        num_warps = num_stages = num_ctas = None
+    num_splits = utils.num_splits_heuristic(
+        seqlen_q=qhead_per_kvhead,
+        seqlen_k=seqlen_k,
+        num_SMs=num_SMs,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        max_split_blocks=utils.max_split_blocks_from_window_sizes(window_sizes, TILE_N)
+        if is_local
+        else None,
+    )
+    out_dtype = torch.bfloat16 if is_quant else query.dtype
+    out = (
+        out
+        if out is not None
+        else torch.empty(query.shape, dtype=out_dtype, device=device)
+    )
+    lse = (
+        lse
+        if lse is not None
+        else torch.empty((batch_size, num_heads_q), dtype=torch.float32, device=device)
+    )
+    out_partial = torch.empty(
+        (num_splits, batch_size, num_heads_q, head_dim),
+        dtype=torch.float32,
+        device=device,
+    )
+    lse_partial = torch.empty(
+        (num_splits, batch_size, num_heads_q),
+        dtype=torch.float32,
+        device=device,
+    )
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    grid = launch_grid.get_dec_grid(
+        batch_size=batch_size,
+        num_heads_kv=num_heads_kv,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        alpha,
+        delta,
+        out_partial,
+        lse_partial,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        softmax_threshold,
+        gate_threshold,
+        window_sizes,
+        query.stride(0),
+        query.stride(-2),
+        1,
+        key.stride(0),
+        key.stride(-2),
+        key.stride(-3),
+        value.stride(0),
+        value.stride(-2),
+        value.stride(-3),
+        alpha.stride(0),
+        1,
+        delta.stride(0),
+        delta.stride(-2),
+        out_partial.stride(1),
+        out_partial.stride(-2),
+        1,
+        out_partial.stride(0),
+        lse_partial.stride(1),
+        lse_partial.stride(-1),
+        1,
+        lse_partial.stride(0),
+        window_sizes.stride(0),
+        None,
+        None,
+        None,
+        None,
+        num_splits,
+        seqlen_q=1,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=0,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD_PACKGQA=qhead_per_kvhead,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_LOCAL=is_local,
+        HAS_CU_SEQLENS_Q=False,
+        HAS_CU_SEQLENS_K=False,
+        HAS_SEQUSED_Q=False,
+        HAS_SEQUSED_K=False,
+        IS_LOGSIGMOID_GATE=is_logsigmoid_gate,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name="dec_gated",
+                seqlen_q=1,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_quant=is_quant,
+            )
+    flash_dec_combine._flash_attn_dec_combine(
+        out_partial,
+        lse_partial,
+        out,
+        lse,
+    )
+    return out, lse
+def _flash_gated_attn_varlen_decode(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    alpha: torch.Tensor,
+    delta: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_k: int,
+    softmax_scale: float = None,
+    softmax_threshold: float = None,
+    gate_threshold: float = None,
+    is_logsigmoid_gate: bool = True,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    seqused_k: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    batch_size, num_heads_q, head_dim = query.shape
+    _, num_heads_kv, _ = key.shape
+    seqlen_k = max_seqlen_k
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    softmax_threshold = (
+        softmax_threshold if softmax_threshold is not None else 1 / seqlen_k
+    )
+    gate_threshold = (
+        gate_threshold if gate_threshold is not None else head_dim / seqlen_k
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_dec_inputs(
+            query,
+            key,
+            value,
+            alpha=alpha,
+            delta=delta,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            cu_seqlens_k=cu_seqlens_k,
+            seqused_k=seqused_k,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name="dec_gated",
+        seqlen_q=1,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _dec_gated_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        TILE_M = max(triton.next_power_of_2(qhead_per_kvhead), 16)
+        TILE_N = 128
+        num_warps = num_stages = num_ctas = None
+    num_splits = utils.num_splits_heuristic(
+        seqlen_q=qhead_per_kvhead,
+        seqlen_k=seqlen_k,
+        num_SMs=num_SMs,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        max_split_blocks=utils.max_split_blocks_from_window_sizes(window_sizes, TILE_N)
+        if is_local
+        else None,
+    )
+    out_dtype = torch.bfloat16 if is_quant else query.dtype
+    out = (
+        out
+        if out is not None
+        else torch.empty(query.shape, dtype=out_dtype, device=device)
+    )
+    lse = (
+        lse
+        if lse is not None
+        else torch.empty((batch_size, num_heads_q), dtype=torch.float32, device=device)
+    )
+    out_partial = torch.empty(
+        (num_splits, batch_size, num_heads_q, head_dim),
+        dtype=torch.float32,
+        device=device,
+    )
+    lse_partial = torch.empty(
+        (num_splits, batch_size, num_heads_q),
+        dtype=torch.float32,
+        device=device,
+    )
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    grid = launch_grid.get_dec_grid(
+        batch_size=batch_size,
+        num_heads_kv=num_heads_kv,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        alpha,
+        delta,
+        out_partial,
+        lse_partial,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        softmax_threshold,
+        gate_threshold,
+        window_sizes,
+        query.stride(0),
+        query.stride(-2),
+        1,
+        0,
+        key.stride(-2),
+        key.stride(0),
+        0,
+        value.stride(-2),
+        value.stride(0),
+        alpha.stride(0),
+        1,
+        0,
+        delta.stride(-2),
+        out_partial.stride(1),
+        out_partial.stride(-2),
+        1,
+        out_partial.stride(0),
+        lse_partial.stride(1),
+        lse_partial.stride(-1),
+        1,
+        lse_partial.stride(0),
+        window_sizes.stride(0),
+        None,
+        cu_seqlens_k,
+        None,
+        seqused_k,
+        num_splits,
+        seqlen_q=1,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=0,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD_PACKGQA=qhead_per_kvhead,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_LOCAL=is_local,
+        HAS_CU_SEQLENS_Q=False,
+        HAS_CU_SEQLENS_K=True,
+        HAS_SEQUSED_Q=False,
+        HAS_SEQUSED_K=seqused_k is not None,
+        IS_LOGSIGMOID_GATE=is_logsigmoid_gate,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name="dec_gated",
+                seqlen_q=1,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_quant=is_quant,
+            )
+    flash_dec_combine._flash_attn_dec_combine(
+        out_partial,
+        lse_partial,
+        out,
+        lse,
+    )
+    return out, lse

build/torch-cuda/flash_gated_fwd.py ADDED Viewed

	@@ -0,0 +1,1337 @@

+from typing import Tuple, Optional
+import torch
+import triton
+import triton.language as tl
+from . import (
+    assert_inputs,
+    utils,
+    cache_utils,
+    launch_template,
+    launch_grid,
+    flash_fwd_combine,
+    kernel_repr,
+    autotuner,
+)
+from .scheduler import (
+    AttnFwdGridIndex,
+    AttnFwdConfig,
+    AttnFwdBlockScheduler,
+    AttnFwdPointerScheduler,
+    AttnMaskScheduler,
+    SoftmaxScheduler,
+)
+@triton.jit
+def _fwd_inner_gated_kernel(
+    config: AttnFwdConfig,
+    ptrs_sched: AttnFwdPointerScheduler,
+    mask_sched: AttnMaskScheduler,
+    softmax_sched: SoftmaxScheduler,
+    skip_gate_curr,
+    acc_s,
+    acc_o,
+    q_tile,
+    a_tile,
+    k_ptrs,
+    v_ptrs,
+    d_ptrs,
+    a_max,
+    a_min,
+    gate_max,
+    row_max,
+    row_sum,
+    n_block,
+    n_block_min,
+    IS_MASK: tl.constexpr,
+    MASK_CAUSAL: tl.constexpr,
+    MASK_LOCAL: tl.constexpr,
+    MASK_SINK: tl.constexpr,
+    CHECK_INF: tl.constexpr,
+):
+    skip_gate_next = True
+    if not skip_gate_curr:
+        if IS_MASK:
+            # Apply mask
+            acc_s = mask_sched.apply_mask(
+                acc_s=acc_s,
+                iter_block=n_block,
+                MASK_CAUSAL=MASK_CAUSAL,
+                MASK_LOCAL=MASK_LOCAL,
+                MASK_SINK=MASK_SINK,
+            )
+        # Apply online sparse softmax
+        p, row_max, row_sum, row_scale, skip_softmax = (
+            softmax_sched.online_sparse_softmax(
+                acc_s=acc_s,
+                row_max=row_max,
+                row_sum=row_sum,
+                softmax_threshold_log2=config.softmax_threshold_log2,
+                CHECK_INF=CHECK_INF,
+            )
+        )
+        if not skip_softmax:
+            # Load value tile
+            v_tile = ptrs_sched.load_v(config, v_ptrs, n_block)
+            # Rescale output accumulator
+            acc_o = softmax_sched.rescale_o(
+                acc_o=acc_o,
+                row_scale=row_scale,
+            )
+            # Update output accumulator
+            acc_o += tl.dot(p.to(v_tile.dtype), v_tile)
+    if n_block > n_block_min:
+        # Load next delta tile
+        d_tile = ptrs_sched.load_d(config, d_ptrs, n_block - 1)
+        d_max = tl.max(d_tile)
+        d_min = tl.min(d_tile)
+        # Check if any gates are active for next tile
+        gate_max, skip_gate_next = softmax_sched.online_gate(
+            a_max=a_max,
+            a_min=a_min,
+            d_max=d_max,
+            d_min=d_min,
+            gate_max=gate_max,
+            gate_threshold_log2=config.gate_threshold_log2,
+        )
+        if not skip_gate_next:
+            # Compute attention gates for next tile
+            acc_s = a_tile[:, None] * d_tile[None, :]
+            if config.IS_LOGSIGMOID_GATE:
+                acc_s = softmax_sched.log_sigmoid(
+                    acc_s=acc_s,
+                )
+            # Load next key tile
+            k_tile = ptrs_sched.load_k(config, k_ptrs, n_block - 1)
+            # Compute attention scores for next tile
+            acc_s += tl.dot(q_tile, k_tile.T)
+    return (
+        skip_gate_next,
+        acc_s,
+        acc_o,
+        gate_max,
+        row_max,
+        row_sum,
+    )
+@triton.jit(repr=kernel_repr.fwd_gated_repr)
+def _fwd_gated_kernel(
+    Q,
+    K,
+    V,
+    A,
+    D,
+    Out,
+    Lse,
+    softmax_scale,
+    softmax_threshold,
+    gate_threshold,
+    query_scale,
+    key_scale,
+    value_scale,
+    window_sizes,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_ab,
+    stride_ah,
+    stride_am,
+    stride_db,
+    stride_dh,
+    stride_dn,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_os,
+    stride_lb,
+    stride_lh,
+    stride_ls,
+    stride_wh,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    seqused_q,
+    seqused_k,
+    num_splits,
+    seqlen_q,
+    seqlen_k,
+    head_dim,
+    SEQLEN_Q_CACHE: tl.constexpr,
+    SEQLEN_K_CACHE: tl.constexpr,
+    QHEAD_PER_KVHEAD: tl.constexpr,
+    PACK_GQA: tl.constexpr,
+    QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+    TILE_K: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    IS_LOCAL: tl.constexpr,
+    IS_SPLIT_KV: tl.constexpr,
+    IS_LOGSIGMOID_GATE: tl.constexpr,
+    IS_ADAPT_GATE: tl.constexpr,
+    HAS_CU_SEQLENS_Q: tl.constexpr,
+    HAS_CU_SEQLENS_K: tl.constexpr,
+    HAS_SEQUSED_Q: tl.constexpr,
+    HAS_SEQUSED_K: tl.constexpr,
+):
+    # Create grid index
+    grid_idx = AttnFwdGridIndex.create(
+        num_splits=num_splits,
+        QHEAD_PER_KVHEAD=QHEAD_PER_KVHEAD,
+        IS_SPLIT_KV=IS_SPLIT_KV,
+        PACK_GQA=PACK_GQA,
+    )
+    # Load window sizes
+    (
+        window_size_sink,
+        window_size_left,
+        window_size_right,
+        window_size_dist,
+    ) = grid_idx.load_window_sizes(
+        window_sizes=window_sizes,
+        stride_wh=stride_wh,
+        IS_LOCAL=IS_LOCAL,
+    )
+    # Create config
+    config = AttnFwdConfig.create(
+        softmax_scale=softmax_scale,
+        softmax_threshold=softmax_threshold,
+        gate_threshold=gate_threshold,
+        query_scale=query_scale,
+        key_scale=key_scale,
+        value_scale=value_scale,
+        m_block=grid_idx.m_block,
+        batch_idx=grid_idx.batch_idx,
+        window_size_sink=window_size_sink,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        window_size_dist=window_size_dist,
+        head_dim=head_dim,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        PACK_GQA=PACK_GQA,
+        QHEAD_PER_KVHEAD_PACKGQA=QHEAD_PER_KVHEAD_PACKGQA,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_CAUSAL=IS_CAUSAL,
+        IS_LOGSIGMOID_GATE=IS_LOGSIGMOID_GATE,
+        IS_ADAPT_GATE=IS_ADAPT_GATE,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+        HAS_SEQUSED_Q=HAS_SEQUSED_Q,
+        HAS_SEQUSED_K=HAS_SEQUSED_K,
+    )
+    # Create pointer scheduler
+    ptrs_sched = AttnFwdPointerScheduler.create(
+        config=config,
+        Q=Q,
+        K=K,
+        V=V,
+        A=A,
+        D=D,
+        Out=Out,
+        Lse=Lse,
+        batch_idx=grid_idx.batch_idx,
+        head_idx=grid_idx.head_idx,
+        head_kv_idx=grid_idx.head_kv_idx,
+        split_idx=grid_idx.split_idx,
+        stride_qb=stride_qb,
+        stride_qh=stride_qh,
+        stride_qm=stride_qm,
+        stride_kb=stride_kb,
+        stride_kh=stride_kh,
+        stride_kn=stride_kn,
+        stride_vb=stride_vb,
+        stride_vh=stride_vh,
+        stride_vn=stride_vn,
+        stride_ab=stride_ab,
+        stride_ah=stride_ah,
+        stride_am=stride_am,
+        stride_db=stride_db,
+        stride_dh=stride_dh,
+        stride_dn=stride_dn,
+        stride_ob=stride_ob,
+        stride_oh=stride_oh,
+        stride_om=stride_om,
+        stride_os=stride_os,
+        stride_lb=stride_lb,
+        stride_lh=stride_lh,
+        stride_ls=stride_ls,
+        IS_SPLIT_KV=IS_SPLIT_KV,
+        IS_GATED=True,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+    )
+    # Create block scheduler
+    block_sched = AttnFwdBlockScheduler.create(
+        config=config,
+        split_idx=grid_idx.split_idx,
+        num_splits=num_splits,
+        IS_CAUSAL=IS_CAUSAL,
+        IS_LOCAL=IS_LOCAL,
+        IS_SPLIT_KV=IS_SPLIT_KV,
+    )
+    # Create mask scheduler
+    mask_sched = AttnMaskScheduler.create(config)
+    # Create softmax scheduler
+    softmax_sched = SoftmaxScheduler.create(config)
+    # Create pointers
+    out_ptrs = ptrs_sched.make_out_ptrs(config)
+    lse_ptrs = ptrs_sched.make_lse_ptrs(config)
+    # Early exit if no n_blocks to process
+    if block_sched.is_empty():
+        ptrs_sched.store_empty(config, out_ptrs, lse_ptrs, IS_SPLIT_KV=IS_SPLIT_KV)
+        return
+    q_ptrs = ptrs_sched.make_q_ptrs(config)
+    k_ptrs = ptrs_sched.make_k_ptrs(config)
+    v_ptrs = ptrs_sched.make_v_ptrs(config)
+    a_ptrs = ptrs_sched.make_a_ptrs(config)
+    d_ptrs = ptrs_sched.make_d_ptrs(config)
+    # Initialize accumulators
+    gate_max = tl.full((), float("-inf"), dtype=tl.float32)
+    row_max = tl.full((TILE_M,), float("-inf"), dtype=tl.float32)
+    row_sum = tl.zeros((TILE_M,), dtype=tl.float32)
+    acc_o = tl.zeros((TILE_M, TILE_K), dtype=tl.float32)
+    # Load query tile
+    q_tile = ptrs_sched.load_q(config, q_ptrs)
+    # Load alpha tile
+    a_tile = ptrs_sched.load_a(config, a_ptrs)
+    a_max = tl.max(a_tile)
+    a_min = tl.min(a_tile)
+    # Load key tile
+    k_tile = ptrs_sched.load_k(config, k_ptrs, block_sched.n_block_max - 1)
+    # Load delta tile
+    d_tile = ptrs_sched.load_d(config, d_ptrs, block_sched.n_block_max - 1)
+    d_max = tl.max(d_tile)
+    d_min = tl.min(d_tile)
+    # Check if any gates are active for current tile
+    gate_max, skip_gate_curr = softmax_sched.online_gate(
+        a_max=a_max,
+        a_min=a_min,
+        d_max=d_max,
+        d_min=d_min,
+        gate_max=gate_max,
+        gate_threshold_log2=config.gate_threshold_log2,
+    )
+    # Initialize skip_gate_curr to False for the first iteration
+    skip_gate_curr = False
+    # Compute attention gates for first tile
+    acc_s = a_tile[:, None] * d_tile[None, :]
+    if IS_LOGSIGMOID_GATE:
+        acc_s = softmax_sched.log_sigmoid(
+            acc_s=acc_s,
+        )
+    # Compute attention scores for first tile
+    acc_s += tl.dot(q_tile, k_tile.T)
+    # Process n_blocks with causal masking
+    if IS_CAUSAL or IS_LOCAL:
+        n_block_max_no_mask = block_sched.n_block_max_no_mask
+        for n_block in tl.range(
+            block_sched.n_block_max - 1, block_sched.n_block_max_no_mask - 1, -1
+        ):
+            (
+                skip_gate_curr,
+                acc_s,
+                acc_o,
+                gate_max,
+                row_max,
+                row_sum,
+            ) = _fwd_inner_gated_kernel(
+                config=config,
+                ptrs_sched=ptrs_sched,
+                mask_sched=mask_sched,
+                softmax_sched=softmax_sched,
+                skip_gate_curr=skip_gate_curr,
+                acc_s=acc_s,
+                acc_o=acc_o,
+                q_tile=q_tile,
+                a_tile=a_tile,
+                k_ptrs=k_ptrs,
+                v_ptrs=v_ptrs,
+                d_ptrs=d_ptrs,
+                a_max=a_max,
+                a_min=a_min,
+                gate_max=gate_max,
+                row_max=row_max,
+                row_sum=row_sum,
+                n_block=n_block,
+                n_block_min=block_sched.n_block_max_no_mask,
+                IS_MASK=True,
+                MASK_CAUSAL=IS_CAUSAL,
+                MASK_LOCAL=True if IS_LOCAL else False,
+                MASK_SINK=False,
+                CHECK_INF=True,
+            )
+    else:
+        # First iteration with seqlen masking
+        n_block = block_sched.n_block_max - 1
+        # Triton aggregate does not support attribute reassignment yet
+        # block_sched.n_block_max_no_mask = n_block
+        n_block_max_no_mask = n_block
+        (
+            skip_gate_curr,
+            acc_s,
+            acc_o,
+            gate_max,
+            row_max,
+            row_sum,
+        ) = _fwd_inner_gated_kernel(
+            config=config,
+            ptrs_sched=ptrs_sched,
+            mask_sched=mask_sched,
+            softmax_sched=softmax_sched,
+            skip_gate_curr=skip_gate_curr,
+            acc_s=acc_s,
+            acc_o=acc_o,
+            q_tile=q_tile,
+            a_tile=a_tile,
+            k_ptrs=k_ptrs,
+            v_ptrs=v_ptrs,
+            d_ptrs=d_ptrs,
+            a_max=a_max,
+            a_min=a_min,
+            gate_max=gate_max,
+            row_max=row_max,
+            row_sum=row_sum,
+            n_block=n_block,
+            n_block_min=n_block,
+            IS_MASK=True,
+            MASK_CAUSAL=False,
+            MASK_LOCAL=False,
+            MASK_SINK=False,
+            CHECK_INF=True,
+        )
+    # Process n_blocks without masking
+    if not IS_LOCAL and n_block_max_no_mask > block_sched.n_block_min:
+        # Load key tile
+        k_tile = ptrs_sched.load_k(config, k_ptrs, n_block_max_no_mask - 1)
+        # Load delta tile
+        d_tile = ptrs_sched.load_d(config, d_ptrs, n_block_max_no_mask - 1)
+        d_max = tl.max(d_tile)
+        d_min = tl.min(d_tile)
+        # Check if any gates are active for current tile
+        gate_max, skip_gate_curr = softmax_sched.online_gate(
+            a_max=a_max,
+            a_min=a_min,
+            d_max=d_max,
+            d_min=d_min,
+            gate_max=gate_max,
+            gate_threshold_log2=config.gate_threshold_log2,
+        )
+        # Compute attention gates
+        acc_s = a_tile[:, None] * d_tile[None, :]
+        if IS_LOGSIGMOID_GATE:
+            acc_s = softmax_sched.log_sigmoid(
+                acc_s=acc_s,
+            )
+        # Compute attention scores
+        acc_s += tl.dot(q_tile, k_tile.T)
+        for n_block in tl.range(
+            n_block_max_no_mask - 1, block_sched.n_block_min - 1, -1
+        ):
+            (
+                skip_gate_curr,
+                acc_s,
+                acc_o,
+                gate_max,
+                row_max,
+                row_sum,
+            ) = _fwd_inner_gated_kernel(
+                config=config,
+                ptrs_sched=ptrs_sched,
+                mask_sched=mask_sched,
+                softmax_sched=softmax_sched,
+                skip_gate_curr=skip_gate_curr,
+                acc_s=acc_s,
+                acc_o=acc_o,
+                q_tile=q_tile,
+                a_tile=a_tile,
+                k_ptrs=k_ptrs,
+                v_ptrs=v_ptrs,
+                d_ptrs=d_ptrs,
+                a_max=a_max,
+                a_min=a_min,
+                gate_max=gate_max,
+                row_max=row_max,
+                row_sum=row_sum,
+                n_block=n_block,
+                n_block_min=block_sched.n_block_min,
+                IS_MASK=False,
+                MASK_CAUSAL=False,
+                MASK_LOCAL=False,
+                MASK_SINK=False,
+                CHECK_INF=False,
+            )
+    if IS_LOCAL:
+        # Process n_blocks with local right masking
+        if block_sched.n_block_window_max > block_sched.n_block_window_max_no_mask:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_max - 1
+            )
+            # Load delta tile
+            d_tile = ptrs_sched.load_d(
+                config, d_ptrs, block_sched.n_block_window_max - 1
+            )
+            d_max = tl.max(d_tile)
+            d_min = tl.min(d_tile)
+            # Check if any gates are active for current tile
+            gate_max, skip_gate_curr = softmax_sched.online_gate(
+                a_max=a_max,
+                a_min=a_min,
+                d_max=d_max,
+                d_min=d_min,
+                gate_max=gate_max,
+                gate_threshold_log2=config.gate_threshold_log2,
+            )
+            # Compute attention gates
+            acc_s = a_tile[:, None] * d_tile[None, :]
+            if IS_LOGSIGMOID_GATE:
+                acc_s = softmax_sched.log_sigmoid(
+                    acc_s=acc_s,
+                )
+            # Compute attention scores
+            acc_s += tl.dot(q_tile, k_tile.T)
+            for n_block in tl.range(
+                block_sched.n_block_window_max - 1,
+                block_sched.n_block_window_max_no_mask - 1,
+                -1,
+            ):
+                (
+                    skip_gate_curr,
+                    acc_s,
+                    acc_o,
+                    gate_max,
+                    row_max,
+                    row_sum,
+                ) = _fwd_inner_gated_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    skip_gate_curr=skip_gate_curr,
+                    acc_s=acc_s,
+                    acc_o=acc_o,
+                    q_tile=q_tile,
+                    a_tile=a_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    d_ptrs=d_ptrs,
+                    a_max=a_max,
+                    a_min=a_min,
+                    gate_max=gate_max,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_max_no_mask,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                    CHECK_INF=True,
+                )
+        # Process n_blocks without masking
+        if (
+            block_sched.n_block_window_max_no_mask
+            > block_sched.n_block_window_min_no_mask
+        ):
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_max_no_mask - 1
+            )
+            # Load delta tile
+            d_tile = ptrs_sched.load_d(
+                config, d_ptrs, block_sched.n_block_window_max_no_mask - 1
+            )
+            d_max = tl.max(d_tile)
+            d_min = tl.min(d_tile)
+            # Check if any gates are active for current tile
+            gate_max, skip_gate_curr = softmax_sched.online_gate(
+                a_max=a_max,
+                a_min=a_min,
+                d_max=d_max,
+                d_min=d_min,
+                gate_max=gate_max,
+                gate_threshold_log2=config.gate_threshold_log2,
+            )
+            # Compute attention gates
+            acc_s = a_tile[:, None] * d_tile[None, :]
+            if IS_LOGSIGMOID_GATE:
+                acc_s = softmax_sched.log_sigmoid(
+                    acc_s=acc_s,
+                )
+            # Compute attention scores
+            acc_s += tl.dot(q_tile, k_tile.T)
+            for n_block in tl.range(
+                block_sched.n_block_window_max_no_mask - 1,
+                block_sched.n_block_window_min_no_mask - 1,
+                -1,
+            ):
+                (
+                    skip_gate_curr,
+                    acc_s,
+                    acc_o,
+                    gate_max,
+                    row_max,
+                    row_sum,
+                ) = _fwd_inner_gated_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    skip_gate_curr=skip_gate_curr,
+                    acc_s=acc_s,
+                    acc_o=acc_o,
+                    q_tile=q_tile,
+                    a_tile=a_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    d_ptrs=d_ptrs,
+                    a_max=a_max,
+                    a_min=a_min,
+                    gate_max=gate_max,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_min_no_mask,
+                    IS_MASK=False,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=False,
+                    MASK_SINK=False,
+                    CHECK_INF=False,
+                )
+        # Process n_blocks with local left masking
+        if block_sched.n_block_window_min_no_mask > block_sched.n_block_window_min:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_min_no_mask - 1
+            )
+            # Load delta tile
+            d_tile = ptrs_sched.load_d(
+                config, d_ptrs, block_sched.n_block_window_min_no_mask - 1
+            )
+            d_max = tl.max(d_tile)
+            d_min = tl.min(d_tile)
+            # Check if any gates are active for current tile
+            gate_max, skip_gate_curr = softmax_sched.online_gate(
+                a_max=a_max,
+                a_min=a_min,
+                d_max=d_max,
+                d_min=d_min,
+                gate_max=gate_max,
+                gate_threshold_log2=config.gate_threshold_log2,
+            )
+            # Compute attention gates
+            acc_s = a_tile[:, None] * d_tile[None, :]
+            if IS_LOGSIGMOID_GATE:
+                acc_s = softmax_sched.log_sigmoid(
+                    acc_s=acc_s,
+                )
+            # Compute attention scores
+            acc_s += tl.dot(q_tile, k_tile.T)
+            for n_block in tl.range(
+                block_sched.n_block_window_min_no_mask - 1,
+                block_sched.n_block_window_min - 1,
+                -1,
+            ):
+                (
+                    skip_gate_curr,
+                    acc_s,
+                    acc_o,
+                    gate_max,
+                    row_max,
+                    row_sum,
+                ) = _fwd_inner_gated_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    skip_gate_curr=skip_gate_curr,
+                    acc_s=acc_s,
+                    acc_o=acc_o,
+                    q_tile=q_tile,
+                    a_tile=a_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    d_ptrs=d_ptrs,
+                    a_max=a_max,
+                    a_min=a_min,
+                    gate_max=gate_max,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_min,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                    CHECK_INF=True,
+                )
+        # Process n_blocks with local sink masking
+        if block_sched.n_block_sink_max > block_sched.n_block_sink_min:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(config, k_ptrs, block_sched.n_block_sink_max - 1)
+            # Load delta tile
+            d_tile = ptrs_sched.load_d(config, d_ptrs, block_sched.n_block_sink_max - 1)
+            d_max = tl.max(d_tile)
+            d_min = tl.min(d_tile)
+            # Check if any gates are active for current tile
+            gate_max, skip_gate_curr = softmax_sched.online_gate(
+                a_max=a_max,
+                a_min=a_min,
+                d_max=d_max,
+                d_min=d_min,
+                gate_max=gate_max,
+                gate_threshold_log2=config.gate_threshold_log2,
+            )
+            # Compute attention gates
+            acc_s = a_tile[:, None] * d_tile[None, :]
+            if IS_LOGSIGMOID_GATE:
+                acc_s = softmax_sched.log_sigmoid(
+                    acc_s=acc_s,
+                )
+            # Compute attention scores
+            acc_s += tl.dot(q_tile, k_tile.T)
+            for n_block in tl.range(
+                block_sched.n_block_sink_max - 1,
+                block_sched.n_block_sink_min - 1,
+                -1,
+            ):
+                (
+                    skip_gate_curr,
+                    acc_s,
+                    acc_o,
+                    gate_max,
+                    row_max,
+                    row_sum,
+                ) = _fwd_inner_gated_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    skip_gate_curr=skip_gate_curr,
+                    acc_s=acc_s,
+                    acc_o=acc_o,
+                    q_tile=q_tile,
+                    a_tile=a_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    d_ptrs=d_ptrs,
+                    a_max=a_max,
+                    a_min=a_min,
+                    gate_max=gate_max,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_sink_min,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=True,
+                    CHECK_INF=True,
+                )
+    # Finalize softmax
+    row_scale, lse_tile = softmax_sched.finalize(
+        row_max=row_max,
+        row_sum=row_sum,
+        IS_LOG2=IS_SPLIT_KV,
+    )
+    # Store LSE
+    ptrs_sched.store_lse(config, lse_ptrs, lse_tile)
+    # Finalize rescale
+    acc_o = softmax_sched.rescale_o(
+        acc_o=acc_o,
+        row_scale=row_scale,
+    )
+    # Store output
+    ptrs_sched.store_out(config, out_ptrs, acc_o, IS_SPLIT_KV=IS_SPLIT_KV)
+_fwd_gated_kernel = cache_utils.wrap_kernel(_fwd_gated_kernel)
+_fwd_gated_kernel_autotuned = None
+def _get_autotuned_kernel():
+    global _fwd_gated_kernel_autotuned
+    if _fwd_gated_kernel_autotuned is None:
+        jit_kernel = _fwd_gated_kernel._kernel
+        autotuned = autotuner.make_fwd_gated_autotuned_kernel(jit_kernel)
+        _fwd_gated_kernel_autotuned = autotuner.AutotunedKernel(autotuned)
+    return _fwd_gated_kernel_autotuned
+def _flash_gated_attn_forward(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    alpha: torch.Tensor,
+    delta: torch.Tensor,
+    is_causal: bool = False,
+    softmax_scale: float = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    softmax_threshold: float = None,
+    gate_threshold: float = None,
+    is_logsigmoid_gate: bool = False,
+    is_adapt_gate: bool = True,
+    is_local: bool = False,
+    is_quant: bool = False,
+    is_split_kv: bool = False,
+    pack_gqa: bool = False,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, float, float, float]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    batch_size, seqlen_q, num_heads_q, head_dim = query.shape
+    _, seqlen_k, num_heads_kv, _ = key.shape
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    softmax_threshold = (
+        softmax_threshold if softmax_threshold is not None else 1 / seqlen_k
+    )
+    gate_threshold = (
+        gate_threshold if gate_threshold is not None else head_dim / seqlen_k
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    qhead_per_kvhead_packgqa = num_heads_q // num_heads_kv if pack_gqa else 1
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_fwd_inputs(
+            query,
+            key,
+            value,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            alpha=alpha,
+            delta=delta,
+            cu_seqlens_q=None,
+            cu_seqlens_k=None,
+            seqused_q=None,
+            seqused_k=None,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    _kernel_name = f"fwd_gated{'_split' if is_split_kv else ''}"
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name=_kernel_name,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_causal=is_causal,
+        pack_gqa=pack_gqa,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _fwd_gated_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        # Placeholder for pre-launch computations
+        TILE_M = TILE_N = 64
+        num_warps = num_stages = num_ctas = None
+    num_splits = (
+        utils.num_splits_heuristic(
+            seqlen_q=seqlen_q * qhead_per_kvhead_packgqa,
+            seqlen_k=seqlen_k,
+            num_SMs=num_SMs,
+            TILE_M=TILE_M,
+            TILE_N=TILE_N,
+            max_split_blocks=utils.max_split_blocks_from_window_sizes(
+                window_sizes, TILE_N
+            )
+            if is_local
+            else None,
+        )
+        if is_split_kv
+        else 1
+    )
+    out_dtype = torch.bfloat16 if is_quant else query.dtype
+    out = out if out is not None else torch.zeros_like(query, dtype=out_dtype)
+    lse = (
+        lse
+        if lse is not None
+        else torch.zeros(
+            (batch_size, num_heads_q, seqlen_q),
+            device=query.device,
+            dtype=torch.float32,
+        )
+    )
+    if is_split_kv:
+        out_partial = torch.empty(
+            (num_splits, batch_size, seqlen_q, num_heads_q, head_dim),
+            dtype=torch.float32,
+            device=query.device,
+        )
+        lse_partial = torch.empty(
+            (num_splits, batch_size, num_heads_q, seqlen_q),
+            dtype=torch.float32,
+            device=query.device,
+        )
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    grid = launch_grid.get_fwd_grid(
+        batch_size=batch_size,
+        seqlen_q=seqlen_q,
+        num_heads_q=num_heads_q,
+        num_heads_kv=num_heads_kv,
+        pack_gqa=pack_gqa,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        alpha,
+        delta,
+        out if not is_split_kv else out_partial,
+        lse if not is_split_kv else lse_partial,
+        softmax_scale,
+        softmax_threshold,
+        gate_threshold,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        query.stride(0),
+        query.stride(-2),
+        query.stride(-3),
+        key.stride(0),
+        key.stride(-2),
+        key.stride(-3),
+        value.stride(0),
+        value.stride(-2),
+        value.stride(-3),
+        alpha.stride(0),
+        alpha.stride(-2),
+        alpha.stride(-1),
+        delta.stride(0),
+        delta.stride(-2),
+        delta.stride(-1),
+        out.stride(0) if not is_split_kv else out_partial.stride(1),
+        out.stride(-2) if not is_split_kv else out_partial.stride(-2),
+        out.stride(-3) if not is_split_kv else out_partial.stride(-3),
+        0 if not is_split_kv else out_partial.stride(0),
+        lse.stride(0) if not is_split_kv else lse_partial.stride(1),
+        lse.stride(-2) if not is_split_kv else lse_partial.stride(-2),
+        0 if not is_split_kv else lse_partial.stride(0),
+        window_sizes.stride(0),
+        None,
+        None,
+        None,
+        None,
+        num_splits,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=seqlen_q // 1024,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD=qhead_per_kvhead,
+        PACK_GQA=pack_gqa,
+        QHEAD_PER_KVHEAD_PACKGQA=qhead_per_kvhead_packgqa,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_CAUSAL=is_causal,
+        IS_LOCAL=is_local,
+        IS_SPLIT_KV=is_split_kv,
+        IS_LOGSIGMOID_GATE=is_logsigmoid_gate,
+        IS_ADAPT_GATE=is_adapt_gate,
+        HAS_CU_SEQLENS_Q=False,
+        HAS_CU_SEQLENS_K=False,
+        HAS_SEQUSED_Q=False,
+        HAS_SEQUSED_K=False,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name=_kernel_name,
+                seqlen_q=seqlen_q,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_causal=is_causal,
+                pack_gqa=pack_gqa,
+                is_quant=is_quant,
+            )
+    if is_split_kv:
+        flash_fwd_combine._flash_attn_fwd_combine(
+            out_partial,
+            lse_partial,
+            out,
+            lse,
+        )
+    return out, lse, softmax_scale, softmax_threshold, gate_threshold
+def _flash_gated_attn_varlen_forward(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    alpha: torch.Tensor,
+    delta: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    is_causal: bool = False,
+    softmax_scale: float = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    softmax_threshold: float = None,
+    gate_threshold: float = None,
+    is_logsigmoid_gate: bool = True,
+    is_adapt_gate: bool = True,
+    is_local: bool = False,
+    is_quant: bool = False,
+    is_split_kv: bool = False,
+    pack_gqa: bool = False,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, float, float, float]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    total_seqlen_q, num_heads_q, head_dim = query.shape
+    _, num_heads_kv, _ = key.shape
+    batch_size = cu_seqlens_q.shape[0] - 1
+    seqlen_q = max_seqlen_q
+    seqlen_k = max_seqlen_k
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    softmax_threshold = (
+        softmax_threshold if softmax_threshold is not None else 1 / seqlen_k
+    )
+    gate_threshold = (
+        gate_threshold if gate_threshold is not None else head_dim / seqlen_k
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    qhead_per_kvhead_packgqa = num_heads_q // num_heads_kv if pack_gqa else 1
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_fwd_inputs(
+            query,
+            key,
+            value,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            alpha=alpha,
+            delta=delta,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    _kernel_name = f"fwd_gated{'_split' if is_split_kv else ''}"
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name=_kernel_name,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_causal=is_causal,
+        pack_gqa=pack_gqa,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _fwd_gated_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        # Placeholder for pre-launch computations
+        TILE_M = TILE_N = 64
+        num_warps = num_stages = num_ctas = None
+    num_splits = (
+        utils.num_splits_heuristic(
+            seqlen_q=seqlen_q * qhead_per_kvhead_packgqa,
+            seqlen_k=seqlen_k,
+            num_SMs=num_SMs,
+            TILE_M=TILE_M,
+            TILE_N=TILE_N,
+            max_split_blocks=utils.max_split_blocks_from_window_sizes(
+                window_sizes, TILE_N
+            )
+            if is_local
+            else None,
+        )
+        if is_split_kv
+        else 1
+    )
+    out_dtype = torch.bfloat16 if is_quant else query.dtype
+    out = out if out is not None else torch.zeros_like(query, dtype=out_dtype)
+    lse = (
+        lse
+        if lse is not None
+        else torch.empty(
+            (num_heads_q, total_seqlen_q),
+            dtype=torch.float32,
+            device=query.device,
+        )
+    )
+    if is_split_kv:
+        out_partial = torch.empty(
+            (num_splits, total_seqlen_q, num_heads_q, head_dim),
+            dtype=torch.float32,
+            device=query.device,
+        )
+        lse_partial = torch.empty(
+            (num_splits, num_heads_q, total_seqlen_q),
+            dtype=torch.float32,
+            device=query.device,
+        )
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    grid = launch_grid.get_fwd_grid(
+        batch_size=batch_size,
+        seqlen_q=seqlen_q,
+        num_heads_q=num_heads_q,
+        num_heads_kv=num_heads_kv,
+        pack_gqa=pack_gqa,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        alpha,
+        delta,
+        out if not is_split_kv else out_partial,
+        lse if not is_split_kv else lse_partial,
+        softmax_scale,
+        softmax_threshold,
+        gate_threshold,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        0,
+        query.stride(-2),
+        query.stride(0),
+        0,
+        key.stride(-2),
+        key.stride(0),
+        0,
+        value.stride(-2),
+        value.stride(0),
+        0,
+        alpha.stride(0),
+        alpha.stride(-1),
+        0,
+        delta.stride(0),
+        delta.stride(-1),
+        0,
+        out.stride(-2) if not is_split_kv else out_partial.stride(-2),
+        out.stride(0) if not is_split_kv else out_partial.stride(-3),
+        0 if not is_split_kv else out_partial.stride(0),
+        0,
+        lse.stride(-2) if not is_split_kv else lse_partial.stride(-2),
+        0 if not is_split_kv else lse_partial.stride(0),
+        window_sizes.stride(0),
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_q,
+        seqused_k,
+        num_splits,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=seqlen_q // 1024,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD=qhead_per_kvhead,
+        PACK_GQA=pack_gqa,
+        QHEAD_PER_KVHEAD_PACKGQA=qhead_per_kvhead_packgqa,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_CAUSAL=is_causal,
+        IS_LOCAL=is_local,
+        IS_SPLIT_KV=is_split_kv,
+        IS_LOGSIGMOID_GATE=is_logsigmoid_gate,
+        IS_ADAPT_GATE=is_adapt_gate,
+        HAS_CU_SEQLENS_Q=True,
+        HAS_CU_SEQLENS_K=True,
+        HAS_SEQUSED_Q=seqused_q is not None,
+        HAS_SEQUSED_K=seqused_k is not None,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name=_kernel_name,
+                seqlen_q=seqlen_q,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_causal=is_causal,
+                pack_gqa=pack_gqa,
+                is_quant=is_quant,
+            )
+    if is_split_kv:
+        flash_fwd_combine._flash_attn_fwd_combine(
+            out_partial,
+            lse_partial,
+            out,
+            lse,
+            cu_seqlens_q=cu_seqlens_q,
+        )
+    return out, lse, softmax_scale, softmax_threshold, gate_threshold

build/torch-cuda/flash_proj_bwd.py ADDED Viewed

	@@ -0,0 +1,21 @@

+import torch
+from triton_kernels.matmul_ogs import matmul_ogs, PrecisionConfig
+def _flash_proj_backward(
+    dout: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    batch_size, seqlen, out_features = dout.shape
+    in_features = input.shape[-1]
+    dy = dout.view(batch_size * seqlen, out_features)
+    x = input.view(batch_size * seqlen, in_features)
+    precision_config = PrecisionConfig(out_dtype=dy.dtype)
+    dx = matmul_ogs(dy, weight, None, precision_config=precision_config)
+    dw = matmul_ogs(dy.T, x, None, precision_config=precision_config)
+    return dx.view(batch_size, seqlen, in_features), dw

build/torch-cuda/flash_proj_fwd.py ADDED Viewed

	@@ -0,0 +1,45 @@

+import torch
+from triton_kernels.matmul_ogs import matmul_ogs, PrecisionConfig, FlexCtx
+from triton_kernels.numerics import InFlexData, OutFlexData
+def _flash_proj_forward(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    out_dtype: torch.dtype = torch.float8_e5m2,
+    prev_scale: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    batch_size, seqlen, in_features = input.shape
+    out_features = weight.shape[0]
+    device = input.device
+    x = input.view(batch_size * seqlen, in_features)
+    if prev_scale is None:
+        expected_scale = torch.tensor([1.0], dtype=torch.float32, device=device)
+    else:
+        expected_scale = prev_scale.view(1).float()
+    actual_scale = torch.zeros(1, dtype=torch.float32, device=device)
+    out_flex = OutFlexData(
+        dtype=out_dtype,
+        expected_scale=expected_scale,
+        actual_scale=actual_scale,
+        checksum_scale=None,
+    )
+    precision_config = PrecisionConfig(
+        flex_ctx=FlexCtx(
+            lhs_data=InFlexData(),
+            rhs_data=InFlexData(),
+            out_data=out_flex,
+        ),
+        out_dtype=out_dtype,
+    )
+    out = matmul_ogs(x, weight.T, None, precision_config=precision_config)
+    descale = precision_config.flex_ctx.out_data.actual_scale.clone()
+    return out.view(batch_size, seqlen, out_features), descale

build/torch-cuda/flash_sparse_attention/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch-cuda/flash_sparse_bwd.py ADDED Viewed

	@@ -0,0 +1,1041 @@

+from typing import Optional, Tuple
+import math
+import torch
+import triton
+import triton.language as tl
+from . import (
+    assert_inputs,
+    utils,
+    cache_utils,
+    launch_template,
+    launch_grid,
+    activations,
+    flash_bwd_preprocess,
+    flash_bwd_postprocess,
+    kernel_repr,
+    autotuner,
+)
+from .scheduler import (
+    AttnBwdGridIndex,
+    AttnBwdConfig,
+    AttnBwdBlockScheduler,
+    AttnBwdPointerScheduler,
+    AttnMaskScheduler,
+)
+@triton.jit
+def _bwd_inner_sparse_kernel(
+    config,
+    ptrs_sched,
+    mask_sched,
+    acc_dk,
+    acc_dv,
+    k_tile,
+    v_tile,
+    q_ptrs,
+    do_ptrs,
+    dq_accum_ptrs,
+    lse_ptrs,
+    dpsum_ptrs,
+    m_block,
+    IS_MASK: tl.constexpr,
+    MASK_CAUSAL: tl.constexpr,
+    MASK_LOCAL: tl.constexpr,
+    MASK_SINK: tl.constexpr,
+):
+    # Load query tile
+    q_tile = ptrs_sched.load_q(config, q_ptrs, m_block)
+    # Compute attention scores
+    acc_s = tl.dot(k_tile, q_tile.T)
+    if IS_MASK:
+        # Apply mask to attention scores
+        acc_s = mask_sched.apply_mask(
+            acc_s=acc_s,
+            iter_block=m_block,
+            MASK_CAUSAL=MASK_CAUSAL,
+            MASK_LOCAL=MASK_LOCAL,
+            MASK_SINK=MASK_SINK,
+        )
+    # Load LSE
+    lse_log2 = ptrs_sched.load_lse(config, lse_ptrs, m_block)
+    # Compute attention weights in log2-domain
+    p_log2 = acc_s * config.softmax_scale_log2 - lse_log2[None, :]
+    # Compute softmax threshold for this m_block
+    softmax_threshold_log2 = config.get_softmax_threshold_log2(
+        m_block=m_block,
+    )
+    # Update skip condition based on threshold
+    skip_softmax = tl.max(p_log2 - softmax_threshold_log2[None, :]) < 0.0
+    if not skip_softmax:
+        # Compute attention weights
+        p = activations.exp2(p_log2).to(q_tile.dtype)
+        # Load output gradients tile
+        do_tile = ptrs_sched.load_do(config, do_ptrs, m_block)
+        # Compute value gradients
+        acc_dv += tl.dot(p, do_tile)
+        # Compute attention weight gradients
+        acc_dp = tl.dot(v_tile, tl.trans(do_tile))
+        # Load dpsum
+        dpsum = ptrs_sched.load_dpsum(config, dpsum_ptrs, m_block)
+        # Compute attention score gradients
+        ds = p * (acc_dp - dpsum[None, :]).to(q_tile.dtype)
+        # Compute query gradients
+        dq = tl.dot(tl.trans(ds), k_tile)
+        # Store query gradients
+        ptrs_sched.store_dq(config, dq_accum_ptrs, m_block, dq)
+        # Compute key gradients
+        acc_dk += tl.dot(ds, q_tile)
+    return acc_dk, acc_dv
+@triton.jit(repr=kernel_repr.bwd_sparse_repr)
+def _bwd_sparse_kernel(
+    Q,
+    K,
+    V,
+    dO,
+    LSELog2,
+    dPsum,
+    dQaccum,
+    dK,
+    dV,
+    softmax_scale,
+    query_scale,
+    key_scale,
+    value_scale,
+    softmax_threshold,
+    window_sizes,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_dob,
+    stride_doh,
+    stride_dom,
+    stride_lb,
+    stride_lh,
+    stride_pb,
+    stride_ph,
+    stride_dqab,
+    stride_dqah,
+    stride_dqam,
+    stride_dkb,
+    stride_dkh,
+    stride_dkn,
+    stride_dks,
+    stride_dvb,
+    stride_dvh,
+    stride_dvn,
+    stride_dvs,
+    stride_wh,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    seqused_q,
+    seqused_k,
+    num_splits,
+    seqlen_q,
+    seqlen_k,
+    head_dim,
+    SEQLEN_Q_CACHE: tl.constexpr,
+    SEQLEN_K_CACHE: tl.constexpr,
+    QHEAD_PER_KVHEAD: tl.constexpr,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+    TILE_K: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    IS_LOCAL: tl.constexpr,
+    IS_SPLIT_QO: tl.constexpr,
+    HAS_CU_SEQLENS_Q: tl.constexpr,
+    HAS_CU_SEQLENS_K: tl.constexpr,
+    HAS_SEQUSED_Q: tl.constexpr,
+    HAS_SEQUSED_K: tl.constexpr,
+):
+    # Create grid index
+    grid_idx = AttnBwdGridIndex.create(
+        num_splits=num_splits,
+        QHEAD_PER_KVHEAD=QHEAD_PER_KVHEAD,
+        IS_SPLIT_QO=IS_SPLIT_QO,
+    )
+    # Load window sizes
+    (
+        window_size_sink,
+        window_size_left,
+        window_size_right,
+        window_size_dist,
+    ) = grid_idx.load_window_sizes(
+        window_sizes=window_sizes,
+        stride_wh=stride_wh,
+        IS_LOCAL=IS_LOCAL,
+    )
+    # Create config
+    config = AttnBwdConfig.create(
+        softmax_scale=softmax_scale,
+        softmax_threshold=softmax_threshold,
+        query_scale=query_scale,
+        key_scale=key_scale,
+        value_scale=value_scale,
+        n_block=grid_idx.n_block,
+        batch_idx=grid_idx.batch_idx,
+        window_size_sink=window_size_sink,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        window_size_dist=window_size_dist,
+        head_dim=head_dim,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        QHEAD_PER_KVHEAD=QHEAD_PER_KVHEAD,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_CAUSAL=IS_CAUSAL,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+        HAS_SEQUSED_Q=HAS_SEQUSED_Q,
+        HAS_SEQUSED_K=HAS_SEQUSED_K,
+    )
+    # Create pointer scheduler
+    ptrs_sched = AttnBwdPointerScheduler.create(
+        config=config,
+        Q=Q,
+        K=K,
+        V=V,
+        dO=dO,
+        LSELog2=LSELog2,
+        dPsum=dPsum,
+        dQaccum=dQaccum,
+        dK=dK,
+        dV=dV,
+        batch_idx=grid_idx.batch_idx,
+        head_idx=grid_idx.head_idx,
+        head_kv_idx=grid_idx.head_kv_idx,
+        split_idx=grid_idx.split_idx,
+        stride_qb=stride_qb,
+        stride_qh=stride_qh,
+        stride_qm=stride_qm,
+        stride_kb=stride_kb,
+        stride_kh=stride_kh,
+        stride_kn=stride_kn,
+        stride_vb=stride_vb,
+        stride_vh=stride_vh,
+        stride_vn=stride_vn,
+        stride_dob=stride_dob,
+        stride_doh=stride_doh,
+        stride_dom=stride_dom,
+        stride_lb=stride_lb,
+        stride_lh=stride_lh,
+        stride_pb=stride_pb,
+        stride_ph=stride_ph,
+        stride_dqab=stride_dqab,
+        stride_dqah=stride_dqah,
+        stride_dqam=stride_dqam,
+        stride_dkb=stride_dkb,
+        stride_dkh=stride_dkh,
+        stride_dkn=stride_dkn,
+        stride_dks=stride_dks,
+        stride_dvb=stride_dvb,
+        stride_dvh=stride_dvh,
+        stride_dvn=stride_dvn,
+        stride_dvs=stride_dvs,
+        IS_SPLIT_QO=IS_SPLIT_QO,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+    )
+    # Create block scheduler
+    block_sched = AttnBwdBlockScheduler.create(
+        config=config,
+        split_idx=grid_idx.split_idx,
+        num_splits=num_splits,
+        IS_CAUSAL=IS_CAUSAL,
+        IS_LOCAL=IS_LOCAL,
+        IS_SPLIT_QO=IS_SPLIT_QO,
+    )
+    # Create mask scheduler
+    mask_sched = AttnMaskScheduler.create(config, SWAP_AB=True)
+    # Early exit if no n_blocks to process
+    if grid_idx.n_block * TILE_N >= config.actual_seqlen_k:
+        return
+    # Create pointers
+    q_ptrs = ptrs_sched.make_q_ptrs(config)
+    k_ptrs = ptrs_sched.make_k_ptrs(config)
+    v_ptrs = ptrs_sched.make_v_ptrs(config)
+    do_ptrs = ptrs_sched.make_do_ptrs(config)
+    lse_ptrs = ptrs_sched.make_lse_ptrs(config)
+    dpsum_ptrs = ptrs_sched.make_dpsum_ptrs(config)
+    dq_accum_ptrs = ptrs_sched.make_dq_accum_ptrs(config)
+    dk_ptrs = ptrs_sched.make_dk_ptrs(config)
+    dv_ptrs = ptrs_sched.make_dv_ptrs(config)
+    # Initialize accumulators
+    acc_dk = tl.zeros((TILE_N, TILE_K), dtype=tl.float32)
+    acc_dv = tl.zeros((TILE_N, TILE_K), dtype=tl.float32)
+    # Load key tile
+    k_tile = ptrs_sched.load_k(config, k_ptrs)
+    # Load value tile
+    v_tile = ptrs_sched.load_v(config, v_ptrs)
+    # Process m_blocks with causal masking
+    if IS_CAUSAL or IS_LOCAL:
+        for m_block in tl.range(
+            block_sched.m_block_min, block_sched.m_block_min_no_mask
+        ):
+            acc_dk, acc_dv = _bwd_inner_sparse_kernel(
+                config=config,
+                ptrs_sched=ptrs_sched,
+                mask_sched=mask_sched,
+                acc_dk=acc_dk,
+                acc_dv=acc_dv,
+                k_tile=k_tile,
+                v_tile=v_tile,
+                q_ptrs=q_ptrs,
+                do_ptrs=do_ptrs,
+                dq_accum_ptrs=dq_accum_ptrs,
+                lse_ptrs=lse_ptrs,
+                dpsum_ptrs=dpsum_ptrs,
+                m_block=m_block,
+                IS_MASK=True,
+                MASK_CAUSAL=IS_CAUSAL,
+                MASK_LOCAL=True if IS_LOCAL else False,
+                MASK_SINK=False,
+            )
+    # Process m_blocks without masking
+    if not IS_LOCAL and block_sched.m_block_min_no_mask < block_sched.m_block_max:
+        for m_block in tl.range(
+            block_sched.m_block_min_no_mask, block_sched.m_block_max
+        ):
+            acc_dk, acc_dv = _bwd_inner_sparse_kernel(
+                config=config,
+                ptrs_sched=ptrs_sched,
+                mask_sched=mask_sched,
+                acc_dk=acc_dk,
+                acc_dv=acc_dv,
+                k_tile=k_tile,
+                v_tile=v_tile,
+                q_ptrs=q_ptrs,
+                do_ptrs=do_ptrs,
+                dq_accum_ptrs=dq_accum_ptrs,
+                lse_ptrs=lse_ptrs,
+                dpsum_ptrs=dpsum_ptrs,
+                m_block=m_block,
+                IS_MASK=False,
+                MASK_CAUSAL=False,
+                MASK_LOCAL=False,
+                MASK_SINK=False,
+            )
+    if IS_LOCAL:
+        # Process m_blocks with local right masking
+        if block_sched.m_block_window_min < block_sched.m_block_window_min_no_mask:
+            for m_block in tl.range(
+                block_sched.m_block_window_min,
+                block_sched.m_block_window_min_no_mask,
+            ):
+                acc_dk, acc_dv = _bwd_inner_sparse_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    acc_dk=acc_dk,
+                    acc_dv=acc_dv,
+                    k_tile=k_tile,
+                    v_tile=v_tile,
+                    q_ptrs=q_ptrs,
+                    do_ptrs=do_ptrs,
+                    dq_accum_ptrs=dq_accum_ptrs,
+                    lse_ptrs=lse_ptrs,
+                    dpsum_ptrs=dpsum_ptrs,
+                    m_block=m_block,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                )
+        # Process m_blocks without masking
+        if (
+            block_sched.m_block_window_min_no_mask
+            < block_sched.m_block_window_max_no_mask
+        ):
+            for m_block in tl.range(
+                block_sched.m_block_window_min_no_mask,
+                block_sched.m_block_window_max_no_mask,
+            ):
+                acc_dk, acc_dv = _bwd_inner_sparse_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    acc_dk=acc_dk,
+                    acc_dv=acc_dv,
+                    k_tile=k_tile,
+                    v_tile=v_tile,
+                    q_ptrs=q_ptrs,
+                    do_ptrs=do_ptrs,
+                    dq_accum_ptrs=dq_accum_ptrs,
+                    lse_ptrs=lse_ptrs,
+                    dpsum_ptrs=dpsum_ptrs,
+                    m_block=m_block,
+                    IS_MASK=False,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=False,
+                    MASK_SINK=False,
+                )
+        # Process m_blocks with local left masking
+        if block_sched.m_block_window_max_no_mask < block_sched.m_block_window_max:
+            for m_block in tl.range(
+                block_sched.m_block_window_max_no_mask,
+                block_sched.m_block_window_max,
+            ):
+                acc_dk, acc_dv = _bwd_inner_sparse_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    acc_dk=acc_dk,
+                    acc_dv=acc_dv,
+                    k_tile=k_tile,
+                    v_tile=v_tile,
+                    q_ptrs=q_ptrs,
+                    do_ptrs=do_ptrs,
+                    dq_accum_ptrs=dq_accum_ptrs,
+                    lse_ptrs=lse_ptrs,
+                    dpsum_ptrs=dpsum_ptrs,
+                    m_block=m_block,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                )
+        # Process m_blocks with local sink masking
+        if block_sched.m_block_sink_max > block_sched.m_block_sink_min:
+            for m_block in tl.range(
+                block_sched.m_block_sink_min,
+                block_sched.m_block_sink_max,
+            ):
+                acc_dk, acc_dv = _bwd_inner_sparse_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    acc_dk=acc_dk,
+                    acc_dv=acc_dv,
+                    k_tile=k_tile,
+                    v_tile=v_tile,
+                    q_ptrs=q_ptrs,
+                    do_ptrs=do_ptrs,
+                    dq_accum_ptrs=dq_accum_ptrs,
+                    lse_ptrs=lse_ptrs,
+                    dpsum_ptrs=dpsum_ptrs,
+                    m_block=m_block,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=True,
+                )
+    # Store value gradients
+    ptrs_sched.store_dv(config, dv_ptrs, acc_dv)
+    # Scale key gradients
+    acc_dk = acc_dk * softmax_scale
+    # Store key gradients
+    ptrs_sched.store_dk(config, dk_ptrs, acc_dk)
+_bwd_sparse_kernel = cache_utils.wrap_kernel(_bwd_sparse_kernel)
+_bwd_sparse_kernel_autotuned = None
+def _get_autotuned_kernel():
+    global _bwd_sparse_kernel_autotuned
+    if _bwd_sparse_kernel_autotuned is None:
+        jit_kernel = _bwd_sparse_kernel._kernel
+        autotuned = autotuner.make_bwd_sparse_autotuned_kernel(jit_kernel)
+        _bwd_sparse_kernel_autotuned = autotuner.AutotunedKernel(autotuned)
+    return _bwd_sparse_kernel_autotuned
+def _flash_sparse_attn_backward(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    dout: torch.Tensor,
+    lse: torch.Tensor,
+    is_causal: bool = False,
+    softmax_scale: float = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    softmax_threshold: float = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    is_split_qo: bool = False,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    batch_size, seqlen_q, num_heads_q, head_dim = query.shape
+    _, seqlen_k, num_heads_kv, _ = key.shape
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    softmax_threshold = (
+        softmax_threshold if softmax_threshold is not None else 1 / seqlen_k
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_bwd_inputs(
+            query,
+            key,
+            value,
+            out,
+            dout,
+            lse,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            cu_seqlens_q=None,
+            cu_seqlens_k=None,
+            seqused_q=None,
+            seqused_k=None,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name="bwd_sparse",
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_causal=is_causal,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _bwd_sparse_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        # Placeholder for pre-launch computations
+        TILE_M = TILE_N = 64
+        num_warps = num_stages = num_ctas = None
+    num_splits = (
+        utils.num_splits_heuristic(
+            seqlen_q=seqlen_k,
+            seqlen_k=seqlen_q,
+            num_SMs=num_SMs,
+            TILE_M=TILE_N,
+            TILE_N=TILE_M,
+            max_split_blocks=utils.max_split_blocks_from_window_sizes(
+                window_sizes, TILE_M
+            )
+            if is_local
+            else None,
+        )
+        if is_split_qo
+        else 1
+    )
+    seqlen_q_rounded = int(math.ceil(seqlen_q / 128) * 128)
+    head_dim_rounded = int(math.ceil(head_dim / 32) * 32)
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    dq = torch.empty_like(query, dtype=query_scale.dtype)
+    dk = torch.empty_like(key, dtype=key_scale.dtype)
+    dv = torch.empty_like(value, dtype=value_scale.dtype)
+    lse_log2 = torch.empty(
+        (batch_size, num_heads_q, seqlen_q_rounded),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dpsum = torch.empty(
+        (batch_size, num_heads_q, seqlen_q_rounded),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dq_accum = torch.empty(
+        (batch_size, num_heads_q, seqlen_q_rounded * head_dim_rounded),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dk_accum = torch.zeros(
+        (num_splits, batch_size, seqlen_k, num_heads_kv, head_dim)
+        if is_split_qo and num_splits > 1
+        else (batch_size, seqlen_k, num_heads_kv, head_dim),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dv_accum = torch.zeros(
+        (num_splits, batch_size, seqlen_k, num_heads_kv, head_dim)
+        if is_split_qo and num_splits > 1
+        else (batch_size, seqlen_k, num_heads_kv, head_dim),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    flash_bwd_preprocess._flash_attn_bwd_preprocess(
+        out=out,
+        dout=dout,
+        dpsum=dpsum,
+        lse=lse,
+        lse_log2=lse_log2,
+        dq_accum=dq_accum,
+        head_dim_rounded=head_dim_rounded,
+        tile_m=TILE_M,
+        tile_k=TILE_K,
+    )
+    grid = launch_grid.get_bwd_grid(
+        seqlen_k=seqlen_k,
+        num_heads_q=num_heads_q,
+        batch_size=batch_size,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        dout,
+        lse_log2,
+        dpsum,
+        dq_accum,
+        dk_accum,
+        dv_accum,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        softmax_threshold,
+        window_sizes,
+        query.stride(0),
+        query.stride(-2),
+        query.stride(-3),
+        key.stride(0),
+        key.stride(-2),
+        key.stride(-3),
+        value.stride(0),
+        value.stride(-2),
+        value.stride(-3),
+        dout.stride(0),
+        dout.stride(-2),
+        dout.stride(-3),
+        lse_log2.stride(0),
+        lse_log2.stride(-2),
+        dpsum.stride(0),
+        dpsum.stride(-2),
+        dq_accum.stride(0),
+        dq_accum.stride(1),
+        head_dim_rounded,
+        dk_accum.stride(-4) if is_split_qo and num_splits > 1 else dk_accum.stride(0),
+        dk_accum.stride(-2),
+        dk_accum.stride(-3),
+        dk_accum.stride(0) if is_split_qo and num_splits > 1 else 0,
+        dv_accum.stride(-4) if is_split_qo and num_splits > 1 else dv_accum.stride(0),
+        dv_accum.stride(-2),
+        dv_accum.stride(-3),
+        dv_accum.stride(0) if is_split_qo and num_splits > 1 else 0,
+        window_sizes.stride(0),
+        None,
+        None,
+        None,
+        None,
+        num_splits,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=seqlen_q // 1024,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD=qhead_per_kvhead,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_CAUSAL=is_causal,
+        IS_LOCAL=is_local,
+        IS_SPLIT_QO=is_split_qo and num_splits > 1,
+        HAS_CU_SEQLENS_Q=False,
+        HAS_CU_SEQLENS_K=False,
+        HAS_SEQUSED_Q=False,
+        HAS_SEQUSED_K=False,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name="bwd_sparse",
+                seqlen_q=seqlen_q,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_causal=is_causal,
+                is_quant=is_quant,
+            )
+    flash_bwd_postprocess._flash_attn_bwd_postprocess(
+        dq_accum=dq_accum,
+        dq=dq,
+        scale=softmax_scale,
+        head_dim_rounded=head_dim_rounded,
+        tile_m=TILE_M,
+        tile_k=TILE_K,
+    )
+    if is_split_qo and num_splits > 1:
+        dk_accum = dk_accum.sum(dim=0)
+        dv_accum = dv_accum.sum(dim=0)
+    dk.copy_(dk_accum)
+    dv.copy_(dv_accum)
+    return dq, dk, dv
+def _flash_sparse_attn_varlen_backward(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    dout: torch.Tensor,
+    lse: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    is_causal: bool = False,
+    softmax_scale: float = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    softmax_threshold: float = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    is_split_qo: bool = False,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    total_q, num_heads_q, head_dim = query.shape
+    total_k, num_heads_kv, _ = key.shape
+    batch_size = cu_seqlens_q.shape[0] - 1
+    seqlen_q = max_seqlen_q
+    seqlen_k = max_seqlen_k
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    softmax_threshold = (
+        softmax_threshold if softmax_threshold is not None else 1 / seqlen_k
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_bwd_inputs(
+            query,
+            key,
+            value,
+            out,
+            dout,
+            lse,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name="bwd_sparse",
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_causal=is_causal,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _bwd_sparse_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        # Placeholder for pre-launch computations
+        TILE_M = TILE_N = 64
+        num_warps = num_stages = num_ctas = None
+    num_splits = (
+        utils.num_splits_heuristic(
+            seqlen_q=seqlen_k,
+            seqlen_k=seqlen_q,
+            num_SMs=num_SMs,
+            TILE_M=TILE_N,
+            TILE_N=TILE_M,
+            max_split_blocks=utils.max_split_blocks_from_window_sizes(
+                window_sizes, TILE_M
+            )
+            if is_local
+            else None,
+        )
+        if is_split_qo
+        else 1
+    )
+    total_q_rounded_padded = int(math.ceil((total_q + batch_size * 128) / 128) * 128)
+    head_dim_rounded = int(math.ceil(head_dim / 32) * 32)
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    dq = torch.empty_like(query, dtype=query_scale.dtype)
+    dk = torch.empty_like(key, dtype=key_scale.dtype)
+    dv = torch.empty_like(value, dtype=value_scale.dtype)
+    lse_log2 = torch.empty(
+        num_heads_q,
+        total_q_rounded_padded,
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dpsum = torch.empty(
+        num_heads_q,
+        total_q_rounded_padded,
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dq_accum = torch.empty(
+        num_heads_q,
+        total_q_rounded_padded * head_dim_rounded,
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dk_accum = torch.zeros(
+        (num_splits, total_k, num_heads_kv, head_dim)
+        if is_split_qo and num_splits > 1
+        else (total_k, num_heads_kv, head_dim),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    dv_accum = torch.zeros(
+        (num_splits, total_k, num_heads_kv, head_dim)
+        if is_split_qo and num_splits > 1
+        else (total_k, num_heads_kv, head_dim),
+        dtype=torch.float32,
+        device=query.device,
+    )
+    flash_bwd_preprocess._flash_attn_bwd_preprocess(
+        out=out,
+        dout=dout,
+        dpsum=dpsum,
+        lse=lse,
+        lse_log2=lse_log2,
+        dq_accum=dq_accum,
+        head_dim_rounded=head_dim_rounded,
+        cu_seqlens_q=cu_seqlens_q,
+        seqused_q=seqused_q,
+        max_seqlen_q=max_seqlen_q,
+        tile_m=TILE_M,
+        tile_k=TILE_K,
+    )
+    grid = launch_grid.get_bwd_grid(
+        seqlen_k=seqlen_k,
+        num_heads_q=num_heads_q,
+        batch_size=batch_size,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        dout,
+        lse_log2,
+        dpsum,
+        dq_accum,
+        dk_accum,
+        dv_accum,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        softmax_threshold,
+        window_sizes,
+        0,
+        query.stride(-2),
+        query.stride(0),
+        0,
+        key.stride(-2),
+        key.stride(0),
+        0,
+        value.stride(-2),
+        value.stride(0),
+        0,
+        dout.stride(-2),
+        dout.stride(0),
+        0,
+        lse_log2.stride(0),
+        0,
+        dpsum.stride(0),
+        0,
+        dq_accum.stride(0),
+        head_dim_rounded,
+        0,
+        dk_accum.stride(-2),
+        dk_accum.stride(-3),
+        dk_accum.stride(0) if is_split_qo and num_splits > 1 else 0,
+        0,
+        dv_accum.stride(-2),
+        dv_accum.stride(-3),
+        dv_accum.stride(0) if is_split_qo and num_splits > 1 else 0,
+        window_sizes.stride(0),
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_q,
+        seqused_k,
+        num_splits,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=seqlen_q // 1024,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD=qhead_per_kvhead,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_CAUSAL=is_causal,
+        IS_LOCAL=is_local,
+        IS_SPLIT_QO=is_split_qo and num_splits > 1,
+        HAS_CU_SEQLENS_Q=True,
+        HAS_CU_SEQLENS_K=True,
+        HAS_SEQUSED_Q=seqused_q is not None,
+        HAS_SEQUSED_K=seqused_k is not None,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name="bwd_sparse",
+                seqlen_q=seqlen_q,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_causal=is_causal,
+                is_quant=is_quant,
+            )
+    flash_bwd_postprocess._flash_attn_bwd_postprocess(
+        dq_accum=dq_accum,
+        dq=dq,
+        scale=softmax_scale,
+        head_dim_rounded=head_dim_rounded,
+        cu_seqlens_q=cu_seqlens_q,
+        seqused_q=seqused_q,
+        max_seqlen_q=max_seqlen_q,
+        tile_m=TILE_M,
+        tile_k=TILE_K,
+    )
+    if is_split_qo and num_splits > 1:
+        dk_accum = dk_accum.sum(dim=0)
+        dv_accum = dv_accum.sum(dim=0)
+    dk.copy_(dk_accum)
+    dv.copy_(dv_accum)
+    return dq, dk, dv

build/torch-cuda/flash_sparse_dec.py ADDED Viewed

	@@ -0,0 +1,871 @@

+from typing import Tuple, Optional
+import torch
+import triton
+import triton.language as tl
+from . import (
+    assert_inputs,
+    utils,
+    cache_utils,
+    launch_template,
+    launch_grid,
+    flash_dec_combine,
+    kernel_repr,
+    autotuner,
+)
+from .scheduler import (
+    AttnDecGridIndex,
+    AttnDecConfig,
+    AttnDecBlockScheduler,
+    AttnDecPointerScheduler,
+    AttnMaskScheduler,
+    SoftmaxScheduler,
+)
+@triton.jit
+def _dec_inner_sparse_kernel(
+    config: AttnDecConfig,
+    ptrs_sched: AttnDecPointerScheduler,
+    mask_sched: AttnMaskScheduler,
+    softmax_sched: SoftmaxScheduler,
+    q_tile,
+    k_tile,
+    k_ptrs,
+    v_ptrs,
+    acc_o,
+    row_max,
+    row_sum,
+    n_block,
+    n_block_min,
+    IS_MASK: tl.constexpr,
+    MASK_LOCAL: tl.constexpr,
+    MASK_SINK: tl.constexpr,
+    CHECK_INF: tl.constexpr,
+):
+    # Compute attention scores
+    acc_s = tl.dot(q_tile, k_tile.T)
+    # Prefetch next key tile
+    if n_block > n_block_min:
+        # Load next key tile
+        k_tile = ptrs_sched.load_k(config, k_ptrs, n_block - 1)
+    if IS_MASK:
+        # Apply mask to attention scores
+        acc_s = mask_sched.apply_mask(
+            acc_s=acc_s,
+            iter_block=n_block,
+            MASK_LOCAL=MASK_LOCAL,
+            MASK_SINK=MASK_SINK,
+        )
+    # Apply online sparse softmax
+    p, row_max, row_sum, row_scale, skip_softmax = softmax_sched.online_sparse_softmax(
+        acc_s=acc_s,
+        row_max=row_max,
+        row_sum=row_sum,
+        softmax_threshold_log2=config.softmax_threshold_log2,
+        CHECK_INF=CHECK_INF,
+    )
+    if not skip_softmax:
+        # Load value tile
+        v_tile = ptrs_sched.load_v(config, v_ptrs, n_block)
+        # Rescale output accumulator
+        acc_o = softmax_sched.rescale_o(
+            acc_o=acc_o,
+            row_scale=row_scale,
+        )
+        # Update output accumulator
+        acc_o += tl.dot(p.to(v_tile.dtype), v_tile)
+    return k_tile, acc_o, row_max, row_sum
+@triton.jit(repr=kernel_repr.dec_sparse_repr)
+def _dec_sparse_kernel(
+    Q,
+    K,
+    V,
+    Out,
+    Lse,
+    softmax_scale,
+    query_scale,
+    key_scale,
+    value_scale,
+    softmax_threshold,
+    window_sizes,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_os,
+    stride_lb,
+    stride_lh,
+    stride_lm,
+    stride_ls,
+    stride_wh,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    seqused_q,
+    seqused_k,
+    num_splits,
+    seqlen_q,
+    seqlen_k,
+    head_dim,
+    SEQLEN_Q_CACHE: tl.constexpr,
+    SEQLEN_K_CACHE: tl.constexpr,
+    QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+    TILE_K: tl.constexpr,
+    IS_LOCAL: tl.constexpr,
+    HAS_CU_SEQLENS_Q: tl.constexpr,
+    HAS_CU_SEQLENS_K: tl.constexpr,
+    HAS_SEQUSED_Q: tl.constexpr,
+    HAS_SEQUSED_K: tl.constexpr,
+):
+    # Create grid index
+    grid_idx = AttnDecGridIndex.create(
+        num_splits=num_splits,
+    )
+    # Load window sizes
+    (
+        window_size_sink,
+        window_size_left,
+        window_size_right,
+        window_size_dist,
+    ) = grid_idx.load_window_sizes(
+        window_sizes=window_sizes,
+        stride_wh=stride_wh,
+        IS_LOCAL=IS_LOCAL,
+    )
+    # Create config
+    config = AttnDecConfig.create(
+        softmax_scale=softmax_scale,
+        softmax_threshold=softmax_threshold,
+        query_scale=query_scale,
+        key_scale=key_scale,
+        value_scale=value_scale,
+        batch_idx=grid_idx.batch_idx,
+        window_size_sink=window_size_sink,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        window_size_dist=window_size_dist,
+        head_dim=head_dim,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        QHEAD_PER_KVHEAD_PACKGQA=QHEAD_PER_KVHEAD_PACKGQA,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+        HAS_SEQUSED_Q=HAS_SEQUSED_Q,
+        HAS_SEQUSED_K=HAS_SEQUSED_K,
+    )
+    # Create pointer scheduler
+    ptrs_sched = AttnDecPointerScheduler.create(
+        config=config,
+        Q=Q,
+        K=K,
+        V=V,
+        Out=Out,
+        Lse=Lse,
+        batch_idx=grid_idx.batch_idx,
+        head_idx=grid_idx.head_idx,
+        head_kv_idx=grid_idx.head_kv_idx,
+        split_idx=grid_idx.split_idx,
+        stride_qb=stride_qb,
+        stride_qh=stride_qh,
+        stride_qm=stride_qm,
+        stride_kb=stride_kb,
+        stride_kh=stride_kh,
+        stride_kn=stride_kn,
+        stride_vb=stride_vb,
+        stride_vh=stride_vh,
+        stride_vn=stride_vn,
+        stride_ob=stride_ob,
+        stride_oh=stride_oh,
+        stride_om=stride_om,
+        stride_os=stride_os,
+        stride_lb=stride_lb,
+        stride_lh=stride_lh,
+        stride_lm=stride_lm,
+        stride_ls=stride_ls,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+    )
+    # Create block scheduler
+    block_sched = AttnDecBlockScheduler.create(
+        config=config,
+        split_idx=grid_idx.split_idx,
+        num_splits=num_splits,
+        IS_LOCAL=IS_LOCAL,
+    )
+    # Create mask scheduler
+    mask_sched = AttnMaskScheduler.create(config)
+    # Create softmax scheduler
+    softmax_sched = SoftmaxScheduler.create(config)
+    # Create pointers
+    out_ptrs = ptrs_sched.make_out_ptrs(config)
+    lse_ptrs = ptrs_sched.make_lse_ptrs(config)
+    # Early exit if no n_blocks to process
+    if block_sched.is_empty():
+        ptrs_sched.store_empty(config, out_ptrs, lse_ptrs, Out)
+        return
+    q_ptrs = ptrs_sched.make_q_ptrs(config)
+    k_ptrs = ptrs_sched.make_k_ptrs(config)
+    v_ptrs = ptrs_sched.make_v_ptrs(config)
+    # Initialize accumulators
+    row_max = tl.full((TILE_M,), float("-inf"), dtype=tl.float32)
+    row_sum = tl.zeros((TILE_M,), dtype=tl.float32)
+    acc_o = tl.zeros((TILE_M, TILE_K), dtype=tl.float32)
+    # Load query tile
+    q_tile = ptrs_sched.load_q(config, q_ptrs)
+    # Load key tile
+    k_tile = ptrs_sched.load_k(config, k_ptrs, block_sched.n_block_max - 1)
+    # Process n_blocks with seqlen masking
+    for n_block in tl.range(
+        block_sched.n_block_max - 1, block_sched.n_block_max_no_mask - 1, -1
+    ):
+        k_tile, acc_o, row_max, row_sum = _dec_inner_sparse_kernel(
+            config=config,
+            ptrs_sched=ptrs_sched,
+            mask_sched=mask_sched,
+            softmax_sched=softmax_sched,
+            q_tile=q_tile,
+            k_tile=k_tile,
+            k_ptrs=k_ptrs,
+            v_ptrs=v_ptrs,
+            acc_o=acc_o,
+            row_max=row_max,
+            row_sum=row_sum,
+            n_block=n_block,
+            n_block_min=block_sched.n_block_max_no_mask,
+            IS_MASK=True,
+            MASK_LOCAL=True if IS_LOCAL else False,
+            MASK_SINK=False,
+            CHECK_INF=True,
+        )
+    # Process n_blocks without masking
+    if not IS_LOCAL and block_sched.n_block_max_no_mask > block_sched.n_block_min:
+        # Load key tile
+        k_tile = ptrs_sched.load_k(config, k_ptrs, block_sched.n_block_max_no_mask - 1)
+        for n_block in tl.range(
+            block_sched.n_block_max_no_mask - 1, block_sched.n_block_min - 1, -1
+        ):
+            k_tile, acc_o, row_max, row_sum = _dec_inner_sparse_kernel(
+                config=config,
+                ptrs_sched=ptrs_sched,
+                mask_sched=mask_sched,
+                softmax_sched=softmax_sched,
+                q_tile=q_tile,
+                k_tile=k_tile,
+                k_ptrs=k_ptrs,
+                v_ptrs=v_ptrs,
+                acc_o=acc_o,
+                row_max=row_max,
+                row_sum=row_sum,
+                n_block=n_block,
+                n_block_min=block_sched.n_block_min,
+                IS_MASK=False,
+                MASK_LOCAL=False,
+                MASK_SINK=False,
+                CHECK_INF=False,
+            )
+    if IS_LOCAL:
+        # Process n_blocks with local right masking
+        if block_sched.n_block_window_max > block_sched.n_block_window_max_no_mask:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_max - 1
+            )
+            for n_block in tl.range(
+                block_sched.n_block_window_max - 1,
+                block_sched.n_block_window_max_no_mask - 1,
+                -1,
+            ):
+                k_tile, acc_o, row_max, row_sum = _dec_inner_sparse_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    q_tile=q_tile,
+                    k_tile=k_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    acc_o=acc_o,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_max_no_mask,
+                    IS_MASK=True,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                    CHECK_INF=True,
+                )
+        # Process n_blocks without masking
+        if (
+            block_sched.n_block_window_max_no_mask
+            > block_sched.n_block_window_min_no_mask
+        ):
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_max_no_mask - 1
+            )
+            for n_block in tl.range(
+                block_sched.n_block_window_max_no_mask - 1,
+                block_sched.n_block_window_min_no_mask - 1,
+                -1,
+            ):
+                k_tile, acc_o, row_max, row_sum = _dec_inner_sparse_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    q_tile=q_tile,
+                    k_tile=k_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    acc_o=acc_o,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_min_no_mask,
+                    IS_MASK=False,
+                    MASK_LOCAL=False,
+                    MASK_SINK=False,
+                    CHECK_INF=False,
+                )
+        # Process n_blocks with local left masking
+        if block_sched.n_block_window_min_no_mask > block_sched.n_block_window_min:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_min_no_mask - 1
+            )
+            for n_block in tl.range(
+                block_sched.n_block_window_min_no_mask - 1,
+                block_sched.n_block_window_min - 1,
+                -1,
+            ):
+                k_tile, acc_o, row_max, row_sum = _dec_inner_sparse_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    q_tile=q_tile,
+                    k_tile=k_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    acc_o=acc_o,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_min,
+                    IS_MASK=True,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                    CHECK_INF=True,
+                )
+        # Process n_blocks with local sink masking
+        if block_sched.n_block_sink_max > block_sched.n_block_sink_min:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(config, k_ptrs, block_sched.n_block_sink_max - 1)
+            for n_block in tl.range(
+                block_sched.n_block_sink_max - 1,
+                block_sched.n_block_sink_min - 1,
+                -1,
+            ):
+                k_tile, acc_o, row_max, row_sum = _dec_inner_sparse_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    q_tile=q_tile,
+                    k_tile=k_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    acc_o=acc_o,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_sink_min,
+                    IS_MASK=True,
+                    MASK_LOCAL=True,
+                    MASK_SINK=True,
+                    CHECK_INF=True,
+                )
+    # Finalize softmax
+    row_scale, lse_tile = softmax_sched.finalize(
+        row_max=row_max,
+        row_sum=row_sum,
+        IS_LOG2=True,
+    )
+    # Store LSE
+    ptrs_sched.store_lse(config, lse_ptrs, lse_tile)
+    # Finalize rescale
+    acc_o = softmax_sched.rescale_o(
+        acc_o=acc_o,
+        row_scale=row_scale,
+    )
+    # Store output
+    ptrs_sched.store_out(config, out_ptrs, acc_o)
+_dec_sparse_kernel = cache_utils.wrap_kernel(_dec_sparse_kernel)
+_dec_sparse_kernel_autotuned = None
+def _get_autotuned_kernel():
+    global _dec_sparse_kernel_autotuned
+    if _dec_sparse_kernel_autotuned is None:
+        jit_kernel = _dec_sparse_kernel._kernel
+        autotuned = autotuner.make_dec_sparse_autotuned_kernel(jit_kernel)
+        _dec_sparse_kernel_autotuned = autotuner.AutotunedKernel(autotuned)
+    return _dec_sparse_kernel_autotuned
+def _flash_sparse_attn_decode(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    softmax_scale: float = None,
+    softmax_threshold: float = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    batch_size, num_heads_q, head_dim = query.shape
+    _, seqlen_k, num_heads_kv, _ = key.shape
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    softmax_threshold = (
+        softmax_threshold if softmax_threshold is not None else 1 / seqlen_k
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_dec_inputs(
+            query,
+            key,
+            value,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            cu_seqlens_k=None,
+            seqused_k=None,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name="dec_sparse",
+        seqlen_q=1,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _dec_sparse_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        TILE_M = max(triton.next_power_of_2(qhead_per_kvhead), 16)
+        TILE_N = 128
+        num_warps = num_stages = num_ctas = None
+    num_splits = utils.num_splits_heuristic(
+        seqlen_q=qhead_per_kvhead,
+        seqlen_k=seqlen_k,
+        num_SMs=num_SMs,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        max_split_blocks=utils.max_split_blocks_from_window_sizes(window_sizes, TILE_N)
+        if is_local
+        else None,
+    )
+    out_dtype = torch.bfloat16 if is_quant else query.dtype
+    out = (
+        out
+        if out is not None
+        else torch.empty(query.shape, dtype=out_dtype, device=device)
+    )
+    lse = (
+        lse
+        if lse is not None
+        else torch.empty((batch_size, num_heads_q), dtype=torch.float32, device=device)
+    )
+    out_partial = torch.empty(
+        (num_splits, batch_size, num_heads_q, head_dim),
+        dtype=torch.float32,
+        device=device,
+    )
+    lse_partial = torch.empty(
+        (num_splits, batch_size, num_heads_q),
+        dtype=torch.float32,
+        device=device,
+    )
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    grid = launch_grid.get_dec_grid(
+        batch_size=batch_size,
+        num_heads_kv=num_heads_kv,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        out_partial,
+        lse_partial,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        softmax_threshold,
+        window_sizes,
+        query.stride(0),
+        query.stride(-2),
+        1,
+        key.stride(0),
+        key.stride(-2),
+        key.stride(-3),
+        value.stride(0),
+        value.stride(-2),
+        value.stride(-3),
+        out_partial.stride(1),
+        out_partial.stride(-2),
+        1,
+        out_partial.stride(0),
+        lse_partial.stride(1),
+        lse_partial.stride(-1),
+        1,
+        lse_partial.stride(0),
+        window_sizes.stride(0),
+        None,
+        None,
+        None,
+        None,
+        num_splits,
+        seqlen_q=1,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=0,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD_PACKGQA=qhead_per_kvhead,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_LOCAL=is_local,
+        HAS_CU_SEQLENS_Q=False,
+        HAS_CU_SEQLENS_K=False,
+        HAS_SEQUSED_Q=False,
+        HAS_SEQUSED_K=False,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name="dec_sparse",
+                seqlen_q=1,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_quant=is_quant,
+            )
+    flash_dec_combine._flash_attn_dec_combine(
+        out_partial,
+        lse_partial,
+        out,
+        lse,
+    )
+    return out, lse
+def _flash_sparse_attn_varlen_decode(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_k: int,
+    softmax_scale: float = None,
+    softmax_threshold: float = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    seqused_k: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    batch_size, num_heads_q, head_dim = query.shape
+    _, num_heads_kv, _ = key.shape
+    seqlen_k = max_seqlen_k
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    softmax_threshold = (
+        softmax_threshold if softmax_threshold is not None else 1 / seqlen_k
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_dec_inputs(
+            query,
+            key,
+            value,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            cu_seqlens_k=cu_seqlens_k,
+            seqused_k=seqused_k,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name="dec_sparse",
+        seqlen_q=1,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _dec_sparse_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        TILE_M = max(triton.next_power_of_2(qhead_per_kvhead), 16)
+        TILE_N = 128
+        num_warps = num_stages = num_ctas = None
+    num_splits = utils.num_splits_heuristic(
+        seqlen_q=qhead_per_kvhead,
+        seqlen_k=seqlen_k,
+        num_SMs=num_SMs,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        max_split_blocks=utils.max_split_blocks_from_window_sizes(window_sizes, TILE_N)
+        if is_local
+        else None,
+    )
+    out_dtype = torch.bfloat16 if is_quant else query.dtype
+    out = (
+        out
+        if out is not None
+        else torch.empty(query.shape, dtype=out_dtype, device=device)
+    )
+    lse = (
+        lse
+        if lse is not None
+        else torch.empty((batch_size, num_heads_q), dtype=torch.float32, device=device)
+    )
+    out_partial = torch.empty(
+        (num_splits, batch_size, num_heads_q, head_dim),
+        dtype=torch.float32,
+        device=device,
+    )
+    lse_partial = torch.empty(
+        (num_splits, batch_size, num_heads_q),
+        dtype=torch.float32,
+        device=device,
+    )
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    grid = launch_grid.get_dec_grid(
+        batch_size=batch_size,
+        num_heads_kv=num_heads_kv,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        out_partial,
+        lse_partial,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        softmax_threshold,
+        window_sizes,
+        query.stride(0),
+        query.stride(-2),
+        1,
+        0,
+        key.stride(-2),
+        key.stride(0),
+        0,
+        value.stride(-2),
+        value.stride(0),
+        out_partial.stride(1),
+        out_partial.stride(-2),
+        1,
+        out_partial.stride(0),
+        lse_partial.stride(1),
+        lse_partial.stride(-1),
+        1,
+        lse_partial.stride(0),
+        window_sizes.stride(0),
+        None,
+        cu_seqlens_k,
+        None,
+        seqused_k,
+        num_splits,
+        seqlen_q=1,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=0,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD_PACKGQA=qhead_per_kvhead,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_LOCAL=is_local,
+        HAS_CU_SEQLENS_Q=False,
+        HAS_CU_SEQLENS_K=True,
+        HAS_SEQUSED_Q=False,
+        HAS_SEQUSED_K=seqused_k is not None,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name="dec_sparse",
+                seqlen_q=1,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_quant=is_quant,
+            )
+    flash_dec_combine._flash_attn_dec_combine(
+        out_partial,
+        lse_partial,
+        out,
+        lse,
+    )
+    return out, lse

build/torch-cuda/flash_sparse_fwd.py ADDED Viewed

	@@ -0,0 +1,974 @@

+from typing import Tuple, Optional
+import torch
+import triton
+import triton.language as tl
+from . import (
+    assert_inputs,
+    utils,
+    cache_utils,
+    launch_template,
+    launch_grid,
+    flash_fwd_combine,
+    kernel_repr,
+    autotuner,
+)
+from .scheduler import (
+    AttnFwdGridIndex,
+    AttnFwdConfig,
+    AttnFwdBlockScheduler,
+    AttnFwdPointerScheduler,
+    AttnMaskScheduler,
+    SoftmaxScheduler,
+)
+@triton.jit
+def _fwd_inner_sparse_kernel(
+    config: AttnFwdConfig,
+    ptrs_sched: AttnFwdPointerScheduler,
+    mask_sched: AttnMaskScheduler,
+    softmax_sched: SoftmaxScheduler,
+    q_tile,
+    k_tile,
+    k_ptrs,
+    v_ptrs,
+    acc_o,
+    row_max,
+    row_sum,
+    n_block,
+    n_block_min,
+    IS_MASK: tl.constexpr,
+    MASK_CAUSAL: tl.constexpr,
+    MASK_LOCAL: tl.constexpr,
+    MASK_SINK: tl.constexpr,
+    CHECK_INF: tl.constexpr,
+):
+    # Compute attention scores
+    acc_s = tl.dot(q_tile, k_tile.T)
+    # Prefetch next key tile
+    if n_block > n_block_min:
+        # Load next key tile
+        k_tile = ptrs_sched.load_k(config, k_ptrs, n_block - 1)
+    if IS_MASK:
+        # Apply mask to attention scores
+        acc_s = mask_sched.apply_mask(
+            acc_s=acc_s,
+            iter_block=n_block,
+            MASK_CAUSAL=MASK_CAUSAL,
+            MASK_LOCAL=MASK_LOCAL,
+            MASK_SINK=MASK_SINK,
+        )
+    # Apply online sparse softmax
+    p, row_max, row_sum, row_scale, skip_softmax = softmax_sched.online_sparse_softmax(
+        acc_s=acc_s,
+        row_max=row_max,
+        row_sum=row_sum,
+        softmax_threshold_log2=config.softmax_threshold_log2,
+        CHECK_INF=CHECK_INF,
+    )
+    if not skip_softmax:
+        # Load value tile
+        v_tile = ptrs_sched.load_v(config, v_ptrs, n_block)
+        # Rescale output accumulator
+        acc_o = softmax_sched.rescale_o(
+            acc_o=acc_o,
+            row_scale=row_scale,
+        )
+        # Update output accumulator
+        acc_o += tl.dot(p.to(v_tile.dtype), v_tile)
+    return k_tile, acc_o, row_max, row_sum
+@triton.jit(repr=kernel_repr.fwd_sparse_repr)
+def _fwd_sparse_kernel(
+    Q,
+    K,
+    V,
+    Out,
+    Lse,
+    softmax_scale,
+    softmax_threshold,
+    query_scale,
+    key_scale,
+    value_scale,
+    window_sizes,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_os,
+    stride_lb,
+    stride_lh,
+    stride_ls,
+    stride_wh,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    seqused_q,
+    seqused_k,
+    num_splits,
+    seqlen_q,
+    seqlen_k,
+    head_dim,
+    SEQLEN_Q_CACHE: tl.constexpr,
+    SEQLEN_K_CACHE: tl.constexpr,
+    QHEAD_PER_KVHEAD: tl.constexpr,
+    PACK_GQA: tl.constexpr,
+    QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+    TILE_K: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    IS_LOCAL: tl.constexpr,
+    IS_SPLIT_KV: tl.constexpr,
+    HAS_CU_SEQLENS_Q: tl.constexpr,
+    HAS_CU_SEQLENS_K: tl.constexpr,
+    HAS_SEQUSED_Q: tl.constexpr,
+    HAS_SEQUSED_K: tl.constexpr,
+):
+    # Create grid index
+    grid_idx = AttnFwdGridIndex.create(
+        num_splits=num_splits,
+        QHEAD_PER_KVHEAD=QHEAD_PER_KVHEAD,
+        IS_SPLIT_KV=IS_SPLIT_KV,
+        PACK_GQA=PACK_GQA,
+    )
+    # Load window sizes
+    (
+        window_size_sink,
+        window_size_left,
+        window_size_right,
+        window_size_dist,
+    ) = grid_idx.load_window_sizes(
+        window_sizes=window_sizes,
+        stride_wh=stride_wh,
+        IS_LOCAL=IS_LOCAL,
+    )
+    # Create config
+    config = AttnFwdConfig.create(
+        softmax_scale=softmax_scale,
+        softmax_threshold=softmax_threshold,
+        query_scale=query_scale,
+        key_scale=key_scale,
+        value_scale=value_scale,
+        m_block=grid_idx.m_block,
+        batch_idx=grid_idx.batch_idx,
+        window_size_sink=window_size_sink,
+        window_size_left=window_size_left,
+        window_size_right=window_size_right,
+        window_size_dist=window_size_dist,
+        head_dim=head_dim,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        PACK_GQA=PACK_GQA,
+        QHEAD_PER_KVHEAD_PACKGQA=QHEAD_PER_KVHEAD_PACKGQA,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+        HAS_SEQUSED_Q=HAS_SEQUSED_Q,
+        HAS_SEQUSED_K=HAS_SEQUSED_K,
+    )
+    # Create pointer scheduler
+    ptrs_sched = AttnFwdPointerScheduler.create(
+        config=config,
+        Q=Q,
+        K=K,
+        V=V,
+        Out=Out,
+        Lse=Lse,
+        batch_idx=grid_idx.batch_idx,
+        head_idx=grid_idx.head_idx,
+        head_kv_idx=grid_idx.head_kv_idx,
+        split_idx=grid_idx.split_idx,
+        stride_qb=stride_qb,
+        stride_qh=stride_qh,
+        stride_qm=stride_qm,
+        stride_kb=stride_kb,
+        stride_kh=stride_kh,
+        stride_kn=stride_kn,
+        stride_vb=stride_vb,
+        stride_vh=stride_vh,
+        stride_vn=stride_vn,
+        stride_ob=stride_ob,
+        stride_oh=stride_oh,
+        stride_om=stride_om,
+        stride_os=stride_os,
+        stride_lb=stride_lb,
+        stride_lh=stride_lh,
+        stride_ls=stride_ls,
+        IS_SPLIT_KV=IS_SPLIT_KV,
+        HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+        HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+    )
+    # Create block scheduler
+    block_sched = AttnFwdBlockScheduler.create(
+        config=config,
+        split_idx=grid_idx.split_idx,
+        num_splits=num_splits,
+        IS_CAUSAL=IS_CAUSAL,
+        IS_LOCAL=IS_LOCAL,
+        IS_SPLIT_KV=IS_SPLIT_KV,
+    )
+    # Create mask scheduler
+    mask_sched = AttnMaskScheduler.create(config)
+    # Create softmax scheduler
+    softmax_sched = SoftmaxScheduler.create(config)
+    # Create pointers
+    out_ptrs = ptrs_sched.make_out_ptrs(config)
+    lse_ptrs = ptrs_sched.make_lse_ptrs(config)
+    # Early exit if no n_blocks to process
+    if block_sched.is_empty():
+        ptrs_sched.store_empty(config, out_ptrs, lse_ptrs, IS_SPLIT_KV=IS_SPLIT_KV)
+        return
+    q_ptrs = ptrs_sched.make_q_ptrs(config)
+    k_ptrs = ptrs_sched.make_k_ptrs(config)
+    v_ptrs = ptrs_sched.make_v_ptrs(config)
+    # Initialize accumulators
+    row_max = tl.full((TILE_M,), float("-inf"), dtype=tl.float32)
+    row_sum = tl.zeros((TILE_M,), dtype=tl.float32)
+    acc_o = tl.zeros((TILE_M, TILE_K), dtype=tl.float32)
+    # Load query tile
+    q_tile = ptrs_sched.load_q(config, q_ptrs)
+    # Load key tile
+    k_tile = ptrs_sched.load_k(config, k_ptrs, block_sched.n_block_max - 1)
+    # Process n_blocks with causal masking
+    if IS_CAUSAL or IS_LOCAL:
+        n_block_max_no_mask = block_sched.n_block_max_no_mask
+        for n_block in tl.range(
+            block_sched.n_block_max - 1, block_sched.n_block_max_no_mask - 1, -1
+        ):
+            k_tile, acc_o, row_max, row_sum = _fwd_inner_sparse_kernel(
+                config=config,
+                ptrs_sched=ptrs_sched,
+                mask_sched=mask_sched,
+                softmax_sched=softmax_sched,
+                q_tile=q_tile,
+                k_tile=k_tile,
+                k_ptrs=k_ptrs,
+                v_ptrs=v_ptrs,
+                acc_o=acc_o,
+                row_max=row_max,
+                row_sum=row_sum,
+                n_block=n_block,
+                n_block_min=block_sched.n_block_max_no_mask,
+                IS_MASK=True,
+                MASK_CAUSAL=IS_CAUSAL,
+                MASK_LOCAL=True if IS_LOCAL else False,
+                MASK_SINK=False,
+                CHECK_INF=True,
+            )
+    else:
+        # First iteration with seqlen masking
+        n_block = block_sched.n_block_max - 1
+        # Triton aggregate does not support attribute reassignment yet
+        # block_sched.n_block_max_no_mask = n_block
+        n_block_max_no_mask = n_block
+        k_tile, acc_o, row_max, row_sum = _fwd_inner_sparse_kernel(
+            config=config,
+            ptrs_sched=ptrs_sched,
+            mask_sched=mask_sched,
+            softmax_sched=softmax_sched,
+            q_tile=q_tile,
+            k_tile=k_tile,
+            k_ptrs=k_ptrs,
+            v_ptrs=v_ptrs,
+            acc_o=acc_o,
+            row_max=row_max,
+            row_sum=row_sum,
+            n_block=n_block,
+            n_block_min=n_block,
+            IS_MASK=True,
+            MASK_CAUSAL=False,
+            MASK_LOCAL=False,
+            MASK_SINK=False,
+            CHECK_INF=True,
+        )
+    # Process n_blocks without masking
+    if not IS_LOCAL and n_block_max_no_mask > block_sched.n_block_min:
+        # Load key tile
+        k_tile = ptrs_sched.load_k(config, k_ptrs, n_block_max_no_mask - 1)
+        for n_block in tl.range(
+            n_block_max_no_mask - 1, block_sched.n_block_min - 1, -1
+        ):
+            k_tile, acc_o, row_max, row_sum = _fwd_inner_sparse_kernel(
+                config=config,
+                ptrs_sched=ptrs_sched,
+                mask_sched=mask_sched,
+                softmax_sched=softmax_sched,
+                q_tile=q_tile,
+                k_tile=k_tile,
+                k_ptrs=k_ptrs,
+                v_ptrs=v_ptrs,
+                acc_o=acc_o,
+                row_max=row_max,
+                row_sum=row_sum,
+                n_block=n_block,
+                n_block_min=block_sched.n_block_min,
+                IS_MASK=False,
+                MASK_CAUSAL=False,
+                MASK_LOCAL=False,
+                MASK_SINK=False,
+                CHECK_INF=False,
+            )
+    if IS_LOCAL:
+        # Process n_blocks with local right masking
+        if block_sched.n_block_window_max > block_sched.n_block_window_max_no_mask:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_max - 1
+            )
+            for n_block in tl.range(
+                block_sched.n_block_window_max - 1,
+                block_sched.n_block_window_max_no_mask - 1,
+                -1,
+            ):
+                k_tile, acc_o, row_max, row_sum = _fwd_inner_sparse_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    q_tile=q_tile,
+                    k_tile=k_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    acc_o=acc_o,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_max_no_mask,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                    CHECK_INF=True,
+                )
+        # Process n_blocks without masking
+        if (
+            block_sched.n_block_window_max_no_mask
+            > block_sched.n_block_window_min_no_mask
+        ):
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_max_no_mask - 1
+            )
+            for n_block in tl.range(
+                block_sched.n_block_window_max_no_mask - 1,
+                block_sched.n_block_window_min_no_mask - 1,
+                -1,
+            ):
+                k_tile, acc_o, row_max, row_sum = _fwd_inner_sparse_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    q_tile=q_tile,
+                    k_tile=k_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    acc_o=acc_o,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_min_no_mask,
+                    IS_MASK=False,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=False,
+                    MASK_SINK=False,
+                    CHECK_INF=False,
+                )
+        # Process n_blocks with local left masking
+        if block_sched.n_block_window_min_no_mask > block_sched.n_block_window_min:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(
+                config, k_ptrs, block_sched.n_block_window_min_no_mask - 1
+            )
+            for n_block in tl.range(
+                block_sched.n_block_window_min_no_mask - 1,
+                block_sched.n_block_window_min - 1,
+                -1,
+            ):
+                k_tile, acc_o, row_max, row_sum = _fwd_inner_sparse_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    q_tile=q_tile,
+                    k_tile=k_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    acc_o=acc_o,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_window_min,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=False,
+                    CHECK_INF=True,
+                )
+        # Process n_blocks with local sink masking
+        if block_sched.n_block_sink_max > block_sched.n_block_sink_min:
+            # Load key tile
+            k_tile = ptrs_sched.load_k(config, k_ptrs, block_sched.n_block_sink_max - 1)
+            for n_block in tl.range(
+                block_sched.n_block_sink_max - 1,
+                block_sched.n_block_sink_min - 1,
+                -1,
+            ):
+                k_tile, acc_o, row_max, row_sum = _fwd_inner_sparse_kernel(
+                    config=config,
+                    ptrs_sched=ptrs_sched,
+                    mask_sched=mask_sched,
+                    softmax_sched=softmax_sched,
+                    q_tile=q_tile,
+                    k_tile=k_tile,
+                    k_ptrs=k_ptrs,
+                    v_ptrs=v_ptrs,
+                    acc_o=acc_o,
+                    row_max=row_max,
+                    row_sum=row_sum,
+                    n_block=n_block,
+                    n_block_min=block_sched.n_block_sink_min,
+                    IS_MASK=True,
+                    MASK_CAUSAL=False,
+                    MASK_LOCAL=True,
+                    MASK_SINK=True,
+                    CHECK_INF=True,
+                )
+    # Finalize softmax
+    row_scale, lse_tile = softmax_sched.finalize(
+        row_max=row_max,
+        row_sum=row_sum,
+        IS_LOG2=IS_SPLIT_KV,
+    )
+    # Store LSE
+    ptrs_sched.store_lse(config, lse_ptrs, lse_tile)
+    # Finalize rescale
+    acc_o = softmax_sched.rescale_o(
+        acc_o=acc_o,
+        row_scale=row_scale,
+    )
+    # Store output
+    ptrs_sched.store_out(config, out_ptrs, acc_o, IS_SPLIT_KV=IS_SPLIT_KV)
+_fwd_sparse_kernel = cache_utils.wrap_kernel(_fwd_sparse_kernel)
+_fwd_sparse_kernel_autotuned = None
+def _get_autotuned_kernel():
+    global _fwd_sparse_kernel_autotuned
+    if _fwd_sparse_kernel_autotuned is None:
+        jit_kernel = _fwd_sparse_kernel._kernel
+        autotuned = autotuner.make_fwd_sparse_autotuned_kernel(jit_kernel)
+        _fwd_sparse_kernel_autotuned = autotuner.AutotunedKernel(autotuned)
+    return _fwd_sparse_kernel_autotuned
+def _flash_sparse_attn_forward(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    is_causal: bool = False,
+    softmax_scale: float = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    softmax_threshold: float = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    is_split_kv: bool = False,
+    pack_gqa: bool = False,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, float, float]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    batch_size, seqlen_q, num_heads_q, head_dim = query.shape
+    _, seqlen_k, num_heads_kv, _ = key.shape
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    softmax_threshold = (
+        softmax_threshold if softmax_threshold is not None else 1 / seqlen_k
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    qhead_per_kvhead_packgqa = num_heads_q // num_heads_kv if pack_gqa else 1
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_fwd_inputs(
+            query,
+            key,
+            value,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            cu_seqlens_q=None,
+            cu_seqlens_k=None,
+            seqused_q=None,
+            seqused_k=None,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    _kernel_name = f"fwd_sparse{'_split' if is_split_kv else ''}"
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name=_kernel_name,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_causal=is_causal,
+        pack_gqa=pack_gqa,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _fwd_sparse_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        # Placeholder for pre-launch computations
+        TILE_M = TILE_N = 64
+        num_warps = num_stages = num_ctas = None
+    num_splits = (
+        utils.num_splits_heuristic(
+            seqlen_q=seqlen_q * qhead_per_kvhead_packgqa,
+            seqlen_k=seqlen_k,
+            num_SMs=num_SMs,
+            TILE_M=TILE_M,
+            TILE_N=TILE_N,
+            max_split_blocks=utils.max_split_blocks_from_window_sizes(
+                window_sizes, TILE_N
+            )
+            if is_local
+            else None,
+        )
+        if is_split_kv
+        else 1
+    )
+    out_dtype = torch.bfloat16 if is_quant else query.dtype
+    out = out if out is not None else torch.empty_like(query, dtype=out_dtype)
+    lse = (
+        lse
+        if lse is not None
+        else torch.empty(
+            (batch_size, num_heads_q, seqlen_q),
+            dtype=torch.float32,
+            device=query.device,
+        )
+    )
+    if is_split_kv:
+        out_partial = torch.empty(
+            (num_splits, batch_size, seqlen_q, num_heads_q, head_dim),
+            dtype=torch.float32,
+            device=query.device,
+        )
+        lse_partial = torch.empty(
+            (num_splits, batch_size, num_heads_q, seqlen_q),
+            dtype=torch.float32,
+            device=query.device,
+        )
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    grid = launch_grid.get_fwd_grid(
+        batch_size=batch_size,
+        seqlen_q=seqlen_q,
+        num_heads_q=num_heads_q,
+        num_heads_kv=num_heads_kv,
+        pack_gqa=pack_gqa,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        out if not is_split_kv else out_partial,
+        lse if not is_split_kv else lse_partial,
+        softmax_scale,
+        softmax_threshold,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        query.stride(0),
+        query.stride(-2),
+        query.stride(-3),
+        key.stride(0),
+        key.stride(-2),
+        key.stride(-3),
+        value.stride(0),
+        value.stride(-2),
+        value.stride(-3),
+        out.stride(0) if not is_split_kv else out_partial.stride(1),
+        out.stride(-2) if not is_split_kv else out_partial.stride(-2),
+        out.stride(-3) if not is_split_kv else out_partial.stride(-3),
+        0 if not is_split_kv else out_partial.stride(0),
+        lse.stride(0) if not is_split_kv else lse_partial.stride(1),
+        lse.stride(-2) if not is_split_kv else lse_partial.stride(-2),
+        0 if not is_split_kv else lse_partial.stride(0),
+        window_sizes.stride(0),
+        None,
+        None,
+        None,
+        None,
+        num_splits,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=seqlen_q // 1024,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD=qhead_per_kvhead,
+        PACK_GQA=pack_gqa,
+        QHEAD_PER_KVHEAD_PACKGQA=qhead_per_kvhead_packgqa,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_CAUSAL=is_causal,
+        IS_LOCAL=is_local,
+        IS_SPLIT_KV=is_split_kv,
+        HAS_CU_SEQLENS_Q=False,
+        HAS_CU_SEQLENS_K=False,
+        HAS_SEQUSED_Q=False,
+        HAS_SEQUSED_K=False,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name=_kernel_name,
+                seqlen_q=seqlen_q,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_causal=is_causal,
+                pack_gqa=pack_gqa,
+                is_quant=is_quant,
+            )
+    if is_split_kv:
+        flash_fwd_combine._flash_attn_fwd_combine(
+            out_partial,
+            lse_partial,
+            out,
+            lse,
+        )
+    return out, lse, softmax_scale, softmax_threshold
+def _flash_sparse_attn_varlen_forward(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    is_causal: bool = False,
+    softmax_scale: float = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    softmax_threshold: float = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    is_split_kv: bool = False,
+    pack_gqa: bool = False,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, float, float]:
+    device = query.device
+    num_SMs = cache_utils.get_device_num_sms(device)
+    total_seqlen_q, num_heads_q, head_dim = query.shape
+    _, num_heads_kv, _ = key.shape
+    batch_size = cu_seqlens_q.shape[0] - 1
+    seqlen_q = max_seqlen_q
+    seqlen_k = max_seqlen_k
+    softmax_scale = (
+        softmax_scale if softmax_scale is not None else 1.0 / (head_dim**0.5)
+    )
+    softmax_threshold = (
+        softmax_threshold if softmax_threshold is not None else 1 / seqlen_k
+    )
+    qhead_per_kvhead = num_heads_q // num_heads_kv
+    qhead_per_kvhead_packgqa = num_heads_q // num_heads_kv if pack_gqa else 1
+    if is_local and window_sizes is None:
+        window_sizes = utils.window_sizes_heuristic(seqlen_k, num_heads_kv, device)
+    elif not is_local:
+        window_sizes = torch.zeros((num_heads_kv, 4), dtype=torch.int32, device=device)
+    if not skip_checks:
+        assert_inputs.assert_fwd_inputs(
+            query,
+            key,
+            value,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim=head_dim,
+            is_quant=is_quant,
+            device=device,
+        )
+    TILE_K = max(triton.next_power_of_2(head_dim), 16)
+    _kernel_name = f"fwd_sparse{'_split' if is_split_kv else ''}"
+    launch_config = launch_template.load_launch_config(
+        device=device,
+        kernel_name=_kernel_name,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        tile_k=TILE_K,
+        is_local=is_local,
+        qhead_per_kvhead=qhead_per_kvhead,
+        is_causal=is_causal,
+        pack_gqa=pack_gqa,
+        is_quant=is_quant,
+    )
+    if launch_config is not None and not is_autotune:
+        kernel = _fwd_sparse_kernel
+        TILE_M, TILE_N, num_warps, num_stages, num_ctas = launch_config
+    else:
+        kernel = _get_autotuned_kernel()
+        # Placeholder for pre-launch computations
+        TILE_M = TILE_N = 64
+        num_warps = num_stages = num_ctas = None
+    num_splits = (
+        utils.num_splits_heuristic(
+            seqlen_q=seqlen_q * qhead_per_kvhead_packgqa,
+            seqlen_k=seqlen_k,
+            num_SMs=num_SMs,
+            TILE_M=TILE_M,
+            TILE_N=TILE_N,
+            max_split_blocks=utils.max_split_blocks_from_window_sizes(
+                window_sizes, TILE_N
+            )
+            if is_local
+            else None,
+        )
+        if is_split_kv
+        else 1
+    )
+    out_dtype = torch.bfloat16 if is_quant else query.dtype
+    out = out if out is not None else torch.empty_like(query, dtype=out_dtype)
+    lse = (
+        lse
+        if lse is not None
+        else torch.empty(
+            (num_heads_q, total_seqlen_q),
+            dtype=torch.float32,
+            device=query.device,
+        )
+    )
+    if is_split_kv:
+        out_partial = torch.empty(
+            (num_splits, total_seqlen_q, num_heads_q, head_dim),
+            dtype=torch.float32,
+            device=query.device,
+        )
+        lse_partial = torch.empty(
+            (num_splits, num_heads_q, total_seqlen_q),
+            dtype=torch.float32,
+            device=query.device,
+        )
+    if not is_quant:
+        query_scale = torch.ones(1, device=device, dtype=query.dtype)
+        key_scale = torch.ones(1, device=device, dtype=query.dtype)
+        value_scale = torch.ones(1, device=device, dtype=query.dtype)
+    grid = launch_grid.get_fwd_grid(
+        batch_size=batch_size,
+        seqlen_q=seqlen_q,
+        num_heads_q=num_heads_q,
+        num_heads_kv=num_heads_kv,
+        pack_gqa=pack_gqa,
+        num_splits=num_splits,
+    )
+    triton.set_allocator(utils.alloc_fn)
+    kernel[grid](
+        query,
+        key,
+        value,
+        out if not is_split_kv else out_partial,
+        lse if not is_split_kv else lse_partial,
+        softmax_scale,
+        softmax_threshold,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        0,
+        query.stride(-2),
+        query.stride(0),
+        0,
+        key.stride(-2),
+        key.stride(0),
+        0,
+        value.stride(-2),
+        value.stride(0),
+        0,
+        out.stride(-2) if not is_split_kv else out_partial.stride(-2),
+        out.stride(0) if not is_split_kv else out_partial.stride(-3),
+        0 if not is_split_kv else out_partial.stride(0),
+        0,
+        lse.stride(-2) if not is_split_kv else lse_partial.stride(-2),
+        0 if not is_split_kv else lse_partial.stride(0),
+        window_sizes.stride(0),
+        cu_seqlens_q,
+        cu_seqlens_k,
+        seqused_q,
+        seqused_k,
+        num_splits,
+        seqlen_q=seqlen_q,
+        seqlen_k=seqlen_k,
+        head_dim=head_dim,
+        SEQLEN_Q_CACHE=seqlen_q // 1024,
+        SEQLEN_K_CACHE=seqlen_k // 1024,
+        QHEAD_PER_KVHEAD=qhead_per_kvhead,
+        PACK_GQA=pack_gqa,
+        QHEAD_PER_KVHEAD_PACKGQA=qhead_per_kvhead_packgqa,
+        TILE_M=TILE_M,
+        TILE_N=TILE_N,
+        TILE_K=TILE_K,
+        IS_CAUSAL=is_causal,
+        IS_LOCAL=is_local,
+        IS_SPLIT_KV=is_split_kv,
+        HAS_CU_SEQLENS_Q=True,
+        HAS_CU_SEQLENS_K=True,
+        HAS_SEQUSED_Q=seqused_q is not None,
+        HAS_SEQUSED_K=seqused_k is not None,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        num_ctas=num_ctas,
+    )
+    if launch_config is None or is_autotune:
+        best = launch_template.extract_best_config(_get_autotuned_kernel())
+        if best is not None:
+            launch_template.store_launch_config(
+                device=device,
+                kernel_name=_kernel_name,
+                seqlen_q=seqlen_q,
+                seqlen_k=seqlen_k,
+                tile_k=TILE_K,
+                config=best,
+                is_local=is_local,
+                qhead_per_kvhead=qhead_per_kvhead,
+                is_causal=is_causal,
+                pack_gqa=pack_gqa,
+                is_quant=is_quant,
+            )
+    if is_split_kv:
+        flash_fwd_combine._flash_attn_fwd_combine(
+            out_partial,
+            lse_partial,
+            out,
+            lse,
+            cu_seqlens_q=cu_seqlens_q,
+        )
+    return out, lse, softmax_scale, softmax_threshold

build/torch-cuda/interface.py ADDED Viewed

	@@ -0,0 +1,1912 @@

+from typing import Optional, Tuple, Union
+import torch
+from .flash_dense_fwd import (
+    _flash_dense_attn_forward,
+    _flash_dense_attn_varlen_forward,
+)
+from .flash_dense_bwd import (
+    _flash_dense_attn_backward,
+    _flash_dense_attn_varlen_backward,
+)
+from .flash_dense_dec import (
+    _flash_dense_attn_decode,
+    _flash_dense_attn_varlen_decode,
+)
+from .flash_sparse_fwd import (
+    _flash_sparse_attn_forward,
+    _flash_sparse_attn_varlen_forward,
+)
+from .flash_sparse_bwd import (
+    _flash_sparse_attn_backward,
+    _flash_sparse_attn_varlen_backward,
+)
+from .flash_sparse_dec import (
+    _flash_sparse_attn_decode,
+    _flash_sparse_attn_varlen_decode,
+)
+from .flash_gated_fwd import (
+    _flash_gated_attn_forward,
+    _flash_gated_attn_varlen_forward,
+)
+from .flash_gated_bwd import (
+    _flash_gated_attn_backward,
+    _flash_gated_attn_varlen_backward,
+)
+from .flash_gated_dec import (
+    _flash_gated_attn_decode,
+    _flash_gated_attn_varlen_decode,
+)
+from .flash_proj_fwd import _flash_proj_forward
+from .flash_proj_bwd import _flash_proj_backward
+from .utils import ensure_contiguous
+from . import quant
+class FlashDenseAttnFunc(torch.autograd.Function):
+    @staticmethod
+    @ensure_contiguous
+    def forward(
+        ctx,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        is_causal: bool = False,
+        softmax_scale: Optional[float] = None,
+        query_scale: Optional[torch.Tensor] = None,
+        key_scale: Optional[torch.Tensor] = None,
+        value_scale: Optional[torch.Tensor] = None,
+        window_sizes: Optional[torch.Tensor] = None,
+        is_local: bool = False,
+        is_quant: bool = False,
+        is_split_kv: bool = False,
+        is_split_qo: bool = False,
+        pack_gqa: bool = False,
+        out: Optional[torch.Tensor] = None,
+        lse: Optional[torch.Tensor] = None,
+        is_autotune: bool = False,
+        skip_checks: bool = False,
+        return_lse: bool = False,
+    ):
+        # Set is_causal to False if sequence length is 1 to avoid unnecessary masking overhead
+        is_causal = False if query.shape[1] == 1 else is_causal
+        if window_sizes is not None:
+            is_local = True
+        if is_quant and (
+            query_scale is None or key_scale is None or value_scale is None
+        ):
+            query, query_scale = quant.quantize_fp8(query)
+            key, key_scale = quant.quantize_fp8(key)
+            value, value_scale = quant.quantize_fp8(value)
+        out, lse, softmax_scale = _flash_dense_attn_forward(
+            query=query,
+            key=key,
+            value=value,
+            is_causal=is_causal,
+            softmax_scale=softmax_scale,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            is_local=is_local,
+            is_quant=is_quant,
+            is_split_kv=is_split_kv,
+            pack_gqa=pack_gqa,
+            out=out,
+            lse=lse,
+            is_autotune=is_autotune,
+            skip_checks=skip_checks,
+        )
+        ctx.save_for_backward(query, key, value, out, lse)
+        ctx.is_causal = is_causal
+        ctx.softmax_scale = softmax_scale
+        ctx.query_scale = query_scale
+        ctx.key_scale = key_scale
+        ctx.value_scale = value_scale
+        ctx.window_sizes = window_sizes
+        ctx.is_local = is_local
+        ctx.is_quant = is_quant
+        ctx.is_split_qo = is_split_qo
+        ctx.is_autotune = is_autotune
+        ctx.skip_checks = skip_checks
+        if return_lse:
+            # LSE gradient is not supported yet
+            ctx.mark_non_differentiable(lse)
+            return out, lse
+        return out
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, dout: torch.Tensor, *args):
+        query, key, value, out, lse = ctx.saved_tensors
+        dq, dk, dv = _flash_dense_attn_backward(
+            query=query,
+            key=key,
+            value=value,
+            out=out,
+            dout=dout,
+            lse=lse,
+            is_causal=ctx.is_causal,
+            softmax_scale=ctx.softmax_scale,
+            query_scale=ctx.query_scale,
+            key_scale=ctx.key_scale,
+            value_scale=ctx.value_scale,
+            window_sizes=ctx.window_sizes,
+            is_local=ctx.is_local,
+            is_quant=ctx.is_quant,
+            is_split_qo=ctx.is_split_qo,
+            is_autotune=ctx.is_autotune,
+            skip_checks=ctx.skip_checks,
+        )
+        return dq, dk, dv, *((None,) * 30)
+class FlashDenseAttnVarlenFunc(torch.autograd.Function):
+    @staticmethod
+    @ensure_contiguous
+    def forward(
+        ctx,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        cu_seqlens_q: torch.Tensor,
+        cu_seqlens_k: torch.Tensor,
+        max_seqlen_q: int,
+        max_seqlen_k: int,
+        is_causal: bool = False,
+        softmax_scale: Optional[float] = None,
+        query_scale: Optional[torch.Tensor] = None,
+        key_scale: Optional[torch.Tensor] = None,
+        value_scale: Optional[torch.Tensor] = None,
+        window_sizes: Optional[torch.Tensor] = None,
+        is_local: bool = False,
+        is_quant: bool = False,
+        is_split_kv: bool = False,
+        is_split_qo: bool = False,
+        pack_gqa: bool = False,
+        seqused_q: Optional[torch.Tensor] = None,
+        seqused_k: Optional[torch.Tensor] = None,
+        out: Optional[torch.Tensor] = None,
+        lse: Optional[torch.Tensor] = None,
+        is_autotune: bool = False,
+        skip_checks: bool = False,
+        return_lse: bool = False,
+    ):
+        # Set is_causal to False if sequence length is 1 to avoid unnecessary masking overhead
+        is_causal = False if max_seqlen_q == 1 else is_causal
+        if window_sizes is not None:
+            is_local = True
+        if is_quant and (
+            query_scale is None or key_scale is None or value_scale is None
+        ):
+            query, query_scale = quant.quantize_fp8(query)
+            key, key_scale = quant.quantize_fp8(key)
+            value, value_scale = quant.quantize_fp8(value)
+        out, lse, softmax_scale = _flash_dense_attn_varlen_forward(
+            query=query,
+            key=key,
+            value=value,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            is_causal=is_causal,
+            softmax_scale=softmax_scale,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            is_local=is_local,
+            is_quant=is_quant,
+            is_split_kv=is_split_kv,
+            pack_gqa=pack_gqa,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+            out=out,
+            lse=lse,
+            is_autotune=is_autotune,
+            skip_checks=skip_checks,
+        )
+        ctx.save_for_backward(
+            query,
+            key,
+            value,
+            out,
+            lse,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqused_q,
+            seqused_k,
+        )
+        ctx.max_seqlen_q = max_seqlen_q
+        ctx.max_seqlen_k = max_seqlen_k
+        ctx.is_causal = is_causal
+        ctx.softmax_scale = softmax_scale
+        ctx.query_scale = query_scale
+        ctx.key_scale = key_scale
+        ctx.value_scale = value_scale
+        ctx.window_sizes = window_sizes
+        ctx.is_local = is_local
+        ctx.is_quant = is_quant
+        ctx.is_split_qo = is_split_qo
+        ctx.is_autotune = is_autotune
+        ctx.skip_checks = skip_checks
+        if return_lse:
+            # LSE gradient is not supported yet
+            ctx.mark_non_differentiable(lse)
+            return out, lse
+        return out
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, dout: torch.Tensor, *args):
+        (
+            query,
+            key,
+            value,
+            out,
+            lse,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqused_q,
+            seqused_k,
+        ) = ctx.saved_tensors
+        dq, dk, dv = _flash_dense_attn_varlen_backward(
+            query=query,
+            key=key,
+            value=value,
+            out=out,
+            dout=dout,
+            lse=lse,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=ctx.max_seqlen_q,
+            max_seqlen_k=ctx.max_seqlen_k,
+            is_causal=ctx.is_causal,
+            softmax_scale=ctx.softmax_scale,
+            query_scale=ctx.query_scale,
+            key_scale=ctx.key_scale,
+            value_scale=ctx.value_scale,
+            window_sizes=ctx.window_sizes,
+            is_local=ctx.is_local,
+            is_quant=ctx.is_quant,
+            is_split_qo=ctx.is_split_qo,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+            is_autotune=ctx.is_autotune,
+            skip_checks=ctx.skip_checks,
+        )
+        return dq, dk, dv, *((None,) * 30)
+class FlashSparseAttnFunc(torch.autograd.Function):
+    @staticmethod
+    @ensure_contiguous
+    def forward(
+        ctx,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        is_causal: bool = False,
+        softmax_scale: Optional[float] = None,
+        query_scale: Optional[torch.Tensor] = None,
+        key_scale: Optional[torch.Tensor] = None,
+        value_scale: Optional[torch.Tensor] = None,
+        window_sizes: Optional[torch.Tensor] = None,
+        softmax_threshold: Optional[float] = None,
+        is_local: bool = False,
+        is_quant: bool = False,
+        is_split_kv: bool = False,
+        is_split_qo: bool = False,
+        pack_gqa: bool = False,
+        out: Optional[torch.Tensor] = None,
+        lse: Optional[torch.Tensor] = None,
+        is_autotune: bool = False,
+        skip_checks: bool = False,
+        return_lse: bool = False,
+    ):
+        # Set is_causal to False if sequence length is 1 to avoid unnecessary masking overhead
+        is_causal = False if query.shape[1] == 1 else is_causal
+        if window_sizes is not None:
+            is_local = True
+        if is_quant and (
+            query_scale is None or key_scale is None or value_scale is None
+        ):
+            query, query_scale = quant.quantize_fp8(query)
+            key, key_scale = quant.quantize_fp8(key)
+            value, value_scale = quant.quantize_fp8(value)
+        out, lse, softmax_scale, softmax_threshold = _flash_sparse_attn_forward(
+            query=query,
+            key=key,
+            value=value,
+            is_causal=is_causal,
+            softmax_scale=softmax_scale,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            softmax_threshold=softmax_threshold,
+            is_local=is_local,
+            is_quant=is_quant,
+            is_split_kv=is_split_kv,
+            pack_gqa=pack_gqa,
+            out=out,
+            lse=lse,
+            is_autotune=is_autotune,
+            skip_checks=skip_checks,
+        )
+        ctx.save_for_backward(query, key, value, out, lse)
+        ctx.is_causal = is_causal
+        ctx.softmax_scale = softmax_scale
+        ctx.query_scale = query_scale
+        ctx.key_scale = key_scale
+        ctx.value_scale = value_scale
+        ctx.window_sizes = window_sizes
+        ctx.softmax_threshold = softmax_threshold
+        ctx.is_local = is_local
+        ctx.is_quant = is_quant
+        ctx.is_split_qo = is_split_qo
+        ctx.is_autotune = is_autotune
+        ctx.skip_checks = skip_checks
+        if return_lse:
+            # LSE gradient is not supported yet
+            ctx.mark_non_differentiable(lse)
+            return out, lse
+        return out
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, dout: torch.Tensor, *args):
+        query, key, value, out, lse = ctx.saved_tensors
+        dq, dk, dv = _flash_sparse_attn_backward(
+            query=query,
+            key=key,
+            value=value,
+            out=out,
+            dout=dout,
+            lse=lse,
+            is_causal=ctx.is_causal,
+            softmax_scale=ctx.softmax_scale,
+            query_scale=ctx.query_scale,
+            key_scale=ctx.key_scale,
+            value_scale=ctx.value_scale,
+            window_sizes=ctx.window_sizes,
+            softmax_threshold=ctx.softmax_threshold,
+            is_local=ctx.is_local,
+            is_quant=ctx.is_quant,
+            is_split_qo=ctx.is_split_qo,
+            is_autotune=ctx.is_autotune,
+            skip_checks=ctx.skip_checks,
+        )
+        return dq, dk, dv, *((None,) * 30)
+class FlashSparseAttnVarlenFunc(torch.autograd.Function):
+    @staticmethod
+    @ensure_contiguous
+    def forward(
+        ctx,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        cu_seqlens_q: torch.Tensor,
+        cu_seqlens_k: torch.Tensor,
+        max_seqlen_q: int,
+        max_seqlen_k: int,
+        is_causal: bool = False,
+        softmax_scale: Optional[float] = None,
+        query_scale: Optional[torch.Tensor] = None,
+        key_scale: Optional[torch.Tensor] = None,
+        value_scale: Optional[torch.Tensor] = None,
+        window_sizes: Optional[torch.Tensor] = None,
+        softmax_threshold: Optional[float] = None,
+        is_local: bool = False,
+        is_quant: bool = False,
+        is_split_kv: bool = False,
+        is_split_qo: bool = False,
+        pack_gqa: bool = False,
+        seqused_q: Optional[torch.Tensor] = None,
+        seqused_k: Optional[torch.Tensor] = None,
+        out: Optional[torch.Tensor] = None,
+        lse: Optional[torch.Tensor] = None,
+        is_autotune: bool = False,
+        skip_checks: bool = False,
+        return_lse: bool = False,
+    ):
+        # Set is_causal to False if sequence length is 1 to avoid unnecessary masking overhead
+        is_causal = False if max_seqlen_q == 1 else is_causal
+        if window_sizes is not None:
+            is_local = True
+        if is_quant and (
+            query_scale is None or key_scale is None or value_scale is None
+        ):
+            query, query_scale = quant.quantize_fp8(query)
+            key, key_scale = quant.quantize_fp8(key)
+            value, value_scale = quant.quantize_fp8(value)
+        out, lse, softmax_scale, softmax_threshold = _flash_sparse_attn_varlen_forward(
+            query=query,
+            key=key,
+            value=value,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            is_causal=is_causal,
+            softmax_scale=softmax_scale,
+            query_scale=query_scale,
+            key_scale=key_scale,
+            value_scale=value_scale,
+            window_sizes=window_sizes,
+            softmax_threshold=softmax_threshold,
+            is_local=is_local,
+            is_quant=is_quant,
+            is_split_kv=is_split_kv,
+            pack_gqa=pack_gqa,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+            out=out,
+            lse=lse,
+            is_autotune=is_autotune,
+            skip_checks=skip_checks,
+        )
+        ctx.save_for_backward(
+            query,
+            key,
+            value,
+            out,
+            lse,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqused_q,
+            seqused_k,
+        )
+        ctx.max_seqlen_q = max_seqlen_q
+        ctx.max_seqlen_k = max_seqlen_k
+        ctx.is_causal = is_causal
+        ctx.softmax_scale = softmax_scale
+        ctx.query_scale = query_scale
+        ctx.key_scale = key_scale
+        ctx.value_scale = value_scale
+        ctx.window_sizes = window_sizes
+        ctx.softmax_threshold = softmax_threshold
+        ctx.is_local = is_local
+        ctx.is_quant = is_quant
+        ctx.is_split_qo = is_split_qo
+        ctx.is_autotune = is_autotune
+        ctx.skip_checks = skip_checks
+        if return_lse:
+            # LSE gradient is not supported yet
+            ctx.mark_non_differentiable(lse)
+            return out, lse
+        return out
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, dout: torch.Tensor, *args):
+        (
+            query,
+            key,
+            value,
+            out,
+            lse,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqused_q,
+            seqused_k,
+        ) = ctx.saved_tensors
+        dq, dk, dv = _flash_sparse_attn_varlen_backward(
+            query=query,
+            key=key,
+            value=value,
+            out=out,
+            dout=dout,
+            lse=lse,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=ctx.max_seqlen_q,
+            max_seqlen_k=ctx.max_seqlen_k,
+            is_causal=ctx.is_causal,
+            softmax_scale=ctx.softmax_scale,
+            query_scale=ctx.query_scale,
+            key_scale=ctx.key_scale,
+            value_scale=ctx.value_scale,
+            window_sizes=ctx.window_sizes,
+            softmax_threshold=ctx.softmax_threshold,
+            is_local=ctx.is_local,
+            is_quant=ctx.is_quant,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+            is_split_qo=ctx.is_split_qo,
+            is_autotune=ctx.is_autotune,
+            skip_checks=ctx.skip_checks,
+        )
+        return dq, dk, dv, *((None,) * 30)
+class FlashGatedAttnFunc(torch.autograd.Function):
+    @staticmethod
+    @ensure_contiguous
+    def forward(
+        ctx,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        alpha: torch.Tensor,
+        delta: torch.Tensor,
+        is_causal: bool = False,
+        softmax_scale: Optional[float] = None,
+        query_scale: Optional[torch.Tensor] = None,
+        key_scale: Optional[torch.Tensor] = None,
+        value_scale: Optional[torch.Tensor] = None,
+        window_sizes: Optional[torch.Tensor] = None,
+        softmax_threshold: Optional[float] = None,
+        gate_threshold: Optional[float] = None,
+        is_logsigmoid_gate: bool = True,
+        is_adapt_gate: bool = True,
+        is_local: bool = False,
+        is_quant: bool = False,
+        is_split_kv: bool = False,
+        is_split_qo: bool = False,
+        pack_gqa: bool = False,
+        out: Optional[torch.Tensor] = None,
+        lse: Optional[torch.Tensor] = None,
+        is_autotune: bool = False,
+        skip_checks: bool = False,
+        return_lse: bool = False,
+    ):
+        # Set is_causal to False if sequence length is 1 to avoid unnecessary masking overhead
+        is_causal = False if query.shape[1] == 1 else is_causal
+        if window_sizes is not None:
+            is_local = True
+        if is_quant and (
+            query_scale is None or key_scale is None or value_scale is None
+        ):
+            query, query_scale = quant.quantize_fp8(query)
+            key, key_scale = quant.quantize_fp8(key)
+            value, value_scale = quant.quantize_fp8(value)
+        # Transpose for make_tensor_descriptor
+        alpha = alpha.transpose(-2, -1).contiguous()
+        delta = delta.transpose(-2, -1).contiguous()
+        out, lse, softmax_scale, softmax_threshold, gate_threshold = (
+            _flash_gated_attn_forward(
+                query=query,
+                key=key,
+                value=value,
+                alpha=alpha,
+                delta=delta,
+                is_causal=is_causal,
+                softmax_scale=softmax_scale,
+                query_scale=query_scale,
+                key_scale=key_scale,
+                value_scale=value_scale,
+                window_sizes=window_sizes,
+                softmax_threshold=softmax_threshold,
+                gate_threshold=gate_threshold,
+                is_logsigmoid_gate=is_logsigmoid_gate,
+                is_adapt_gate=is_adapt_gate,
+                is_local=is_local,
+                is_quant=is_quant,
+                is_split_kv=is_split_kv,
+                pack_gqa=pack_gqa,
+                out=out,
+                lse=lse,
+                is_autotune=is_autotune,
+                skip_checks=skip_checks,
+            )
+        )
+        ctx.save_for_backward(query, key, value, alpha, delta, out, lse)
+        ctx.is_causal = is_causal
+        ctx.softmax_scale = softmax_scale
+        ctx.query_scale = query_scale
+        ctx.key_scale = key_scale
+        ctx.value_scale = value_scale
+        ctx.window_sizes = window_sizes
+        ctx.softmax_threshold = softmax_threshold
+        ctx.gate_threshold = gate_threshold
+        ctx.is_logsigmoid_gate = is_logsigmoid_gate
+        ctx.is_adapt_gate = is_adapt_gate
+        ctx.is_local = is_local
+        ctx.is_quant = is_quant
+        ctx.is_split_qo = is_split_qo
+        ctx.is_autotune = is_autotune
+        ctx.skip_checks = skip_checks
+        if return_lse:
+            # LSE gradient is not supported yet
+            ctx.mark_non_differentiable(lse)
+            return out, lse
+        return out
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, dout: torch.Tensor, *args):
+        query, key, value, alpha, delta, out, lse = ctx.saved_tensors
+        dq, dk, dv, da, dd = _flash_gated_attn_backward(
+            query=query,
+            key=key,
+            value=value,
+            alpha=alpha,
+            delta=delta,
+            out=out,
+            dout=dout,
+            lse=lse,
+            is_causal=ctx.is_causal,
+            softmax_scale=ctx.softmax_scale,
+            query_scale=ctx.query_scale,
+            key_scale=ctx.key_scale,
+            value_scale=ctx.value_scale,
+            window_sizes=ctx.window_sizes,
+            softmax_threshold=ctx.softmax_threshold,
+            gate_threshold=ctx.gate_threshold,
+            is_logsigmoid_gate=ctx.is_logsigmoid_gate,
+            is_adapt_gate=ctx.is_adapt_gate,
+            is_local=ctx.is_local,
+            is_quant=ctx.is_quant,
+            is_split_qo=ctx.is_split_qo,
+            is_autotune=ctx.is_autotune,
+            skip_checks=ctx.skip_checks,
+        )
+        # Transpose back to original layout
+        da = da.transpose(-2, -1).contiguous()
+        dd = dd.transpose(-2, -1).contiguous()
+        return dq, dk, dv, da, dd, *((None,) * 30)
+class FlashGatedAttnVarlenFunc(torch.autograd.Function):
+    @staticmethod
+    @ensure_contiguous
+    def forward(
+        ctx,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        alpha: torch.Tensor,
+        delta: torch.Tensor,
+        cu_seqlens_q: torch.Tensor,
+        cu_seqlens_k: torch.Tensor,
+        max_seqlen_q: int,
+        max_seqlen_k: int,
+        is_causal: bool = False,
+        softmax_scale: Optional[float] = None,
+        query_scale: Optional[torch.Tensor] = None,
+        key_scale: Optional[torch.Tensor] = None,
+        value_scale: Optional[torch.Tensor] = None,
+        window_sizes: Optional[torch.Tensor] = None,
+        softmax_threshold: Optional[float] = None,
+        gate_threshold: Optional[float] = None,
+        is_logsigmoid_gate: bool = True,
+        is_adapt_gate: bool = True,
+        is_local: bool = False,
+        is_quant: bool = False,
+        is_split_kv: bool = False,
+        is_split_qo: bool = False,
+        pack_gqa: bool = False,
+        seqused_q: Optional[torch.Tensor] = None,
+        seqused_k: Optional[torch.Tensor] = None,
+        out: Optional[torch.Tensor] = None,
+        lse: Optional[torch.Tensor] = None,
+        is_autotune: bool = False,
+        skip_checks: bool = False,
+        return_lse: bool = False,
+    ):
+        # Set is_causal to False if sequence length is 1 to avoid unnecessary masking overhead
+        is_causal = False if max_seqlen_q == 1 else is_causal
+        if window_sizes is not None:
+            is_local = True
+        if is_quant and (
+            query_scale is None or key_scale is None or value_scale is None
+        ):
+            query, query_scale = quant.quantize_fp8(query)
+            key, key_scale = quant.quantize_fp8(key)
+            value, value_scale = quant.quantize_fp8(value)
+        # Transpose for make_tensor_descriptor
+        alpha = alpha.transpose(-2, -1).contiguous()
+        delta = delta.transpose(-2, -1).contiguous()
+        out, lse, softmax_scale, softmax_threshold, gate_threshold = (
+            _flash_gated_attn_varlen_forward(
+                query=query,
+                key=key,
+                value=value,
+                alpha=alpha,
+                delta=delta,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_q,
+                max_seqlen_k=max_seqlen_k,
+                is_causal=is_causal,
+                softmax_scale=softmax_scale,
+                query_scale=query_scale,
+                key_scale=key_scale,
+                value_scale=value_scale,
+                window_sizes=window_sizes,
+                softmax_threshold=softmax_threshold,
+                gate_threshold=gate_threshold,
+                is_logsigmoid_gate=is_logsigmoid_gate,
+                is_adapt_gate=is_adapt_gate,
+                is_local=is_local,
+                is_quant=is_quant,
+                is_split_kv=is_split_kv,
+                pack_gqa=pack_gqa,
+                seqused_q=seqused_q,
+                seqused_k=seqused_k,
+                out=out,
+                lse=lse,
+                is_autotune=is_autotune,
+                skip_checks=skip_checks,
+            )
+        )
+        ctx.save_for_backward(
+            query,
+            key,
+            value,
+            alpha,
+            delta,
+            out,
+            lse,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqused_q,
+            seqused_k,
+        )
+        ctx.max_seqlen_q = max_seqlen_q
+        ctx.max_seqlen_k = max_seqlen_k
+        ctx.is_causal = is_causal
+        ctx.softmax_scale = softmax_scale
+        ctx.query_scale = query_scale
+        ctx.key_scale = key_scale
+        ctx.value_scale = value_scale
+        ctx.window_sizes = window_sizes
+        ctx.softmax_threshold = softmax_threshold
+        ctx.gate_threshold = gate_threshold
+        ctx.is_logsigmoid_gate = is_logsigmoid_gate
+        ctx.is_adapt_gate = is_adapt_gate
+        ctx.is_local = is_local
+        ctx.is_quant = is_quant
+        ctx.is_split_qo = is_split_qo
+        ctx.is_autotune = is_autotune
+        ctx.skip_checks = skip_checks
+        if return_lse:
+            # LSE gradient is not supported yet
+            ctx.mark_non_differentiable(lse)
+            return out, lse
+        return out
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, dout: torch.Tensor, *args):
+        (
+            query,
+            key,
+            value,
+            alpha,
+            delta,
+            out,
+            lse,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqused_q,
+            seqused_k,
+        ) = ctx.saved_tensors
+        dq, dk, dv, da, dd = _flash_gated_attn_varlen_backward(
+            query=query,
+            key=key,
+            value=value,
+            alpha=alpha,
+            delta=delta,
+            out=out,
+            dout=dout,
+            lse=lse,
+            softmax_scale=ctx.softmax_scale,
+            query_scale=ctx.query_scale,
+            key_scale=ctx.key_scale,
+            value_scale=ctx.value_scale,
+            window_sizes=ctx.window_sizes,
+            softmax_threshold=ctx.softmax_threshold,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=ctx.max_seqlen_q,
+            max_seqlen_k=ctx.max_seqlen_k,
+            is_causal=ctx.is_causal,
+            gate_threshold=ctx.gate_threshold,
+            is_logsigmoid_gate=ctx.is_logsigmoid_gate,
+            is_adapt_gate=ctx.is_adapt_gate,
+            is_local=ctx.is_local,
+            is_quant=ctx.is_quant,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+            is_split_qo=ctx.is_split_qo,
+            is_autotune=ctx.is_autotune,
+            skip_checks=ctx.skip_checks,
+        )
+        # Transpose back to original layout
+        da = da.transpose(-2, -1).contiguous()
+        dd = dd.transpose(-2, -1).contiguous()
+        return dq, dk, dv, da, dd, *((None,) * 30)
+class FlashProjFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        out_dtype: torch.dtype = torch.float8_e4m3fn,
+        prev_scale: Optional[torch.Tensor] = None,
+    ):
+        out, descale = _flash_proj_forward(
+            input=input,
+            weight=weight,
+            out_dtype=out_dtype,
+            prev_scale=prev_scale,
+        )
+        ctx.save_for_backward(input, weight)
+        return out, descale
+    @staticmethod
+    def backward(ctx, dout: torch.Tensor, _dscale: torch.Tensor):
+        input, weight = ctx.saved_tensors
+        dout = dout.to(input.dtype)
+        dinput, dweight = _flash_proj_backward(dout=dout, input=input, weight=weight)
+        return dinput, dweight, None, None
+def flash_dense_attn_func(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    is_causal: bool = False,
+    softmax_scale: Optional[float] = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    is_split_kv: bool = False,
+    is_split_qo: bool = False,
+    pack_gqa: bool = False,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+    return_lse: bool = False,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Flash dense attention function that computes the attention output and optionally the logsumexp.
+    :param query: Query tensor of shape [batch_size, seqlen_q, num_heads, head_dim].
+    :param key: Key tensor of shape [batch_size, seqlen_k, num_kv_heads, head_dim].
+    :param value: Value tensor of shape [batch_size, seqlen_k, num_kv_heads, head_dim].
+    :param is_causal: Whether to apply a causal mask.
+    :param softmax_scale: Optional scaling factor for the softmax. If None, defaults to 1/sqrt(head_dim).
+    :param query_scale: Optional per-tensor scale for query dequantization.
+    :param key_scale: Optional per-tensor scale for key dequantization.
+    :param value_scale: Optional per-tensor scale for value dequantization.
+    :param window_sizes: Optional window sizes tensor for flexible local attention. Must be shape [num_kv_heads, 4] with dtype int32, with columns [window_sink, window_left, window_right, window_dist]. If provided, is_local is automatically set to True.
+    :param is_local: Whether to apply a local mask.
+    :param is_quant: Whether the inputs are quantized. If True, query_scale, key_scale, and value_scale must be provided or will be computed from the input tensors.
+    :param is_split_kv: Whether to enable split-KV for forward occupancy.
+    :param is_split_qo: Whether to enable split-QO for backward occupancy.
+    :param pack_gqa: Whether to pack grouped-query attention.
+    :param out: Optional preallocated output tensor with shape [batch_size, seqlen_q, num_heads, head_dim].
+    :param lse: Optional preallocated logsumexp tensor with shape [batch_size, num_heads, seqlen_q].
+    :param is_autotune: Force re-run Triton autotuner and overwrite cache. Default False uses cached config.
+    :param skip_checks: Whether to skip input validation checks for faster performance.
+    :param return_lse: Whether to return the logsumexp tensor for numerical stability analysis. If True, returns a tuple (out, lse). If False, returns only out.
+    :returns: If return_lse is False, returns out with shape [batch_size, seqlen_q, num_heads, head_dim]. If return_lse is True, returns a tuple (out, lse), where lse has shape [batch_size, num_heads, seqlen_q].
+    """
+    return FlashDenseAttnFunc.apply(
+        query,
+        key,
+        value,
+        is_causal,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        is_local,
+        is_quant,
+        is_split_kv,
+        is_split_qo,
+        pack_gqa,
+        out,
+        lse,
+        is_autotune,
+        skip_checks,
+        return_lse,
+    )
+def flash_dense_attn_with_kvcache_func(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    softmax_scale: Optional[float] = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+    return_lse: bool = False,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Flash dense attention function for decoding with KV cache that computes the attention output and optionally the logsumexp.
+    :param query: Query tensor of shape [batch_size, num_heads, head_dim].
+    :param key: Key tensor of shape [batch_size, seqlen_k, num_kv_heads, head_dim].
+    :param value: Value tensor of shape [batch_size, seqlen_k, num_kv_heads, head_dim].
+    :param softmax_scale: Optional scaling factor for the softmax. If None, defaults to 1/sqrt(head_dim).
+    :param query_scale: Optional per-tensor scale for query dequantization.
+    :param key_scale: Optional per-tensor scale for key dequantization.
+    :param value_scale: Optional per-tensor scale for value dequantization.
+    :param window_sizes: Optional window sizes tensor for flexible local attention. Must be shape [num_kv_heads, 4] with dtype int32, with columns [window_sink, window_left, window_right, window_dist]. If provided, is_local is automatically set to True.
+    :param is_local: Whether to apply a local mask.
+    :param is_quant: Whether the inputs are quantized. If True, query_scale, key_scale, and value_scale must be provided for dequantization.
+    :param out: Optional preallocated output tensor with shape [batch_size, num_heads, head_dim].
+    :param lse: Optional preallocated logsumexp tensor with shape [batch_size, num_heads].
+    :param is_autotune: Force re-run Triton autotuner and overwrite cache. Default False uses cached config.
+    :param skip_checks: Whether to skip input validation checks for faster performance.
+    :param return_lse: Whether to return the logsumexp tensor for numerical stability analysis. If True, returns a tuple (out, lse). If False, returns only out.
+    :returns: If return_lse is False, returns out with shape [batch_size, num_heads, head_dim]. If return_lse is True, returns a tuple (out, lse), where lse has shape [batch_size, num_heads].
+    """
+    if is_quant and (query_scale is None or key_scale is None or value_scale is None):
+        query, query_scale = quant.quantize_fp8(query)
+        key, key_scale = quant.quantize_fp8(key)
+        value, value_scale = quant.quantize_fp8(value)
+    if window_sizes is not None:
+        is_local = True
+    out, lse = _flash_dense_attn_decode(
+        query=query,
+        key=key,
+        value=value,
+        softmax_scale=softmax_scale,
+        query_scale=query_scale,
+        key_scale=key_scale,
+        value_scale=value_scale,
+        window_sizes=window_sizes,
+        is_local=is_local,
+        is_quant=is_quant,
+        out=out,
+        lse=lse,
+        is_autotune=is_autotune,
+        skip_checks=skip_checks,
+    )
+    if return_lse:
+        return out, lse
+    return out
+def flash_dense_attn_varlen_func(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    is_causal: bool = False,
+    softmax_scale: Optional[float] = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    is_split_kv: bool = False,
+    is_split_qo: bool = False,
+    pack_gqa: bool = False,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+    return_lse: bool = False,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Flash dense attention function for variable-length sequences that computes the attention output and optionally the logsumexp.
+    :param query: Query tensor of shape [total_seqlen_q, num_heads_q, head_dim].
+    :param key: Key tensor of shape [total_seqlen_k, num_heads_kv, head_dim].
+    :param value: Value tensor of shape [total_seqlen_k, num_heads_kv, head_dim].
+    :param cu_seqlens_q: Cumulative sequence lengths for queries, shape [batch_size + 1].
+    :param cu_seqlens_k: Cumulative sequence lengths for keys/values, shape [batch_size + 1].
+    :param max_seqlen_q: Maximum sequence length for queries.
+    :param max_seqlen_k: Maximum sequence length for keys/values.
+    :param is_causal: Whether to apply a causal mask.
+    :param softmax_scale: Optional scaling factor for the softmax. If None, defaults to 1/sqrt(head_dim).
+    :param query_scale: Optional per-tensor scale for query dequantization.
+    :param key_scale: Optional per-tensor scale for key dequantization.
+    :param value_scale: Optional per-tensor scale for value dequantization.
+    :param window_sizes: Optional window sizes tensor for flexible local attention. Must be shape [num_kv_heads, 4] with dtype int32, with columns [window_sink, window_left, window_right, window_dist]. If provided, is_local is automatically set to True.
+    :param is_local: Whether to apply a local mask.
+    :param is_quant: Whether the inputs are quantized. If True, query_scale, key_scale, and value_scale must be provided or will be computed from the input tensors.
+    :param is_split_kv: Whether to enable split-KV for forward occupancy.
+    :param is_split_qo: Whether to enable split-QO for backward occupancy.
+    :param pack_gqa: Whether to pack grouped-query attention.
+    :param seqused_q: Optional tensor of shape [total_seqlen_q] indicating the actual sequence lengths for queries. If provided, overrides cu_seqlens_q for masking.
+    :param seqused_k: Optional tensor of shape [total_seqlen_k] indicating the actual sequence lengths for keys/values. If provided, overrides cu_seqlens_k for masking.
+    :param out: Optional preallocated output tensor with shape [batch_size, seqlen_q, num_heads, head_dim].
+    :param lse: Optional preallocated logsumexp tensor with shape [batch_size, num_heads, seqlen_q].
+    :param is_autotune: Force re-run Triton autotuner and overwrite cache. Default False uses cached config.
+    :param skip_checks: Whether to skip input validation checks for faster performance.
+    :param return_lse: Whether to return the logsumexp tensor for numerical stability analysis. If True, returns a tuple (out, lse). If False, returns only out.
+    :returns: If return_lse is False, returns out with shape [total_seqlen_q, num_heads_q, head_dim]. If return_lse is True, returns a tuple (out, lse), where lse has shape [total_seqlen_q, num_heads_q].
+    """
+    return FlashDenseAttnVarlenFunc.apply(
+        query,
+        key,
+        value,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        is_causal,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        is_local,
+        is_quant,
+        is_split_kv,
+        is_split_qo,
+        pack_gqa,
+        seqused_q,
+        seqused_k,
+        out,
+        lse,
+        is_autotune,
+        skip_checks,
+        return_lse,
+    )
+def flash_dense_attn_varlen_with_kvcache_func(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_k: int,
+    softmax_scale: Optional[float] = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    seqused_k: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+    return_lse: bool = False,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Flash dense attention function for variable-length decoding with KV cache that computes the attention output and optionally the logsumexp.
+    :param query: Query tensor of shape [batch_size, num_heads_q, head_dim].
+    :param key: Key tensor of shape [total_seqlen_k, num_heads_kv, head_dim].
+    :param value: Value tensor of shape [total_seqlen_k, num_heads_kv, head_dim].
+    :param cu_seqlens_k: Cumulative sequence lengths for keys/values, shape [batch_size + 1].
+    :param max_seqlen_k: Maximum sequence length for keys/values.
+    :param softmax_scale: Optional scaling factor for the softmax. If None, defaults to 1/sqrt(head_dim).
+    :param query_scale: Optional per-tensor scale for query dequantization.
+    :param key_scale: Optional per-tensor scale for key dequantization.
+    :param value_scale: Optional per-tensor scale for value dequantization.
+    :param window_sizes: Optional window sizes tensor for flexible local attention. Must be shape [num_kv_heads, 4] with dtype int32, with columns [window_sink, window_left, window_right, window_dist]. If provided, is_local is automatically set to True.
+    :param is_local: Whether to apply a local mask.
+    :param is_quant: Whether the inputs are quantized. If True, query_scale, key_scale, and value_scale must be provided for dequantization.
+    :param seqused_k: Optional tensor indicating the actual sequence lengths for keys/values.
+    :param out: Optional preallocated output tensor with shape [batch_size, num_heads_q, head_dim].
+    :param lse: Optional preallocated logsumexp tensor with shape [batch_size, num_heads_q].
+    :param is_autotune: Force re-run Triton autotuner and overwrite cache. Default False uses cached config.
+    :param skip_checks: Whether to skip input validation checks for faster performance.
+    :param return_lse: Whether to return the logsumexp tensor for numerical stability analysis. If True, returns a tuple (out, lse). If False, returns only out.
+    :returns: If return_lse is False, returns out with shape [batch_size, num_heads_q, head_dim]. If return_lse is True, returns a tuple (out, lse), where lse has shape [batch_size, num_heads_q].
+    """
+    if is_quant and (query_scale is None or key_scale is None or value_scale is None):
+        query, query_scale = quant.quantize_fp8(query)
+        key, key_scale = quant.quantize_fp8(key)
+        value, value_scale = quant.quantize_fp8(value)
+    if window_sizes is not None:
+        is_local = True
+    out, lse = _flash_dense_attn_varlen_decode(
+        query=query,
+        key=key,
+        value=value,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_k=max_seqlen_k,
+        softmax_scale=softmax_scale,
+        query_scale=query_scale,
+        key_scale=key_scale,
+        value_scale=value_scale,
+        window_sizes=window_sizes,
+        is_local=is_local,
+        is_quant=is_quant,
+        seqused_k=seqused_k,
+        out=out,
+        lse=lse,
+        is_autotune=is_autotune,
+        skip_checks=skip_checks,
+    )
+    if return_lse:
+        return out, lse
+    return out
+def flash_sparse_attn_func(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    is_causal: bool = False,
+    softmax_scale: Optional[float] = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    softmax_threshold: Optional[float] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    is_split_kv: bool = False,
+    is_split_qo: bool = False,
+    pack_gqa: bool = False,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+    return_lse: bool = False,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Flash sparse attention function that computes the attention output and optionally the logsumexp.
+    :param query: Query tensor of shape [batch_size, seqlen_q, num_heads, head_dim].
+    :param key: Key tensor of shape [batch_size, seqlen_k, num_kv_heads, head_dim].
+    :param value: Value tensor of shape [batch_size, seqlen_k, num_kv_heads, head_dim].
+    :param is_causal: Whether to apply a causal mask.
+    :param softmax_scale: Optional scaling factor for the softmax. If None, defaults to 1/sqrt(head_dim).
+    :param query_scale: Optional per-tensor scale for query dequantization.
+    :param key_scale: Optional per-tensor scale for key dequantization.
+    :param value_scale: Optional per-tensor scale for value dequantization.
+    :param window_sizes: Optional window sizes tensor for flexible local attention. Must be shape [num_kv_heads, 4] with dtype int32, with columns [window_sink, window_left, window_right, window_dist]. If provided, is_local is automatically set to True.
+    :param softmax_threshold: Optional threshold for the sparse softmax. If None, defaults to head_dim / seqlen_k.
+    :param is_local: Whether to apply a local mask.
+    :param is_quant: Whether the inputs are quantized. If True, query_scale, key_scale, and value_scale must be provided or will be computed from the input tensors.
+    :param is_split_kv: Whether to enable split-KV for forward ccupancy.
+    :param is_split_qo: Whether to enable split-QO for backward occupancy.
+    :param pack_gqa: Whether to pack grouped-query attention.
+    :param out: Optional preallocated output tensor with shape [batch_size, seqlen_q, num_heads, head_dim].
+    :param lse: Optional preallocated logsumexp tensor with shape [batch_size, num_heads, seqlen_q].
+    :param is_autotune: Force re-run Triton autotuner and overwrite cache. Default False uses cached config.
+    :param skip_checks: Whether to skip input validation checks for faster performance.
+    :param return_lse: Whether to return the logsumexp tensor for numerical stability analysis. If True, returns a tuple (out, lse). If False, returns only out.
+    :returns: If return_lse is False, returns out with shape [batch_size, seqlen_q, num_heads, head_dim]. If return_lse is True, returns a tuple (out, lse), where lse has shape [batch_size, num_heads, seqlen_q].
+    """
+    return FlashSparseAttnFunc.apply(
+        query,
+        key,
+        value,
+        is_causal,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        softmax_threshold,
+        is_local,
+        is_quant,
+        is_split_kv,
+        is_split_qo,
+        pack_gqa,
+        out,
+        lse,
+        is_autotune,
+        skip_checks,
+        return_lse,
+    )
+def flash_sparse_attn_with_kvcache_func(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    softmax_scale: Optional[float] = None,
+    softmax_threshold: Optional[float] = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+    return_lse: bool = False,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Flash sparse attention function for decoding with KV cache that computes the attention output and optionally the logsumexp.
+    :param query: Query tensor of shape [batch_size, num_heads, head_dim].
+    :param key: Key tensor of shape [batch_size, seqlen_k, num_kv_heads, head_dim].
+    :param value: Value tensor of shape [batch_size, seqlen_k, num_kv_heads, head_dim].
+    :param softmax_scale: Optional scaling factor for the softmax. If None, defaults to 1/sqrt(head_dim).
+    :param softmax_threshold: Optional threshold for the sparse softmax. If None, defaults to head_dim / seqlen_k.
+    :param query_scale: Optional per-tensor scale for query dequantization.
+    :param key_scale: Optional per-tensor scale for key dequantization.
+    :param value_scale: Optional per-tensor scale for value dequantization.
+    :param window_sizes: Optional window sizes tensor for flexible local attention. Must be shape [num_kv_heads, 4] with dtype int32, with columns [window_sink, window_left, window_right, window_dist]. If provided, is_local is automatically set to True.
+    :param is_local: Whether to apply a local mask.
+    :param is_quant: Whether the inputs are quantized. If True, query_scale, key_scale, and value_scale must be provided for dequantization.
+    :param out: Optional preallocated output tensor with shape [batch_size, num_heads, head_dim].
+    :param lse: Optional preallocated logsumexp tensor with shape [batch_size, num_heads].
+    :param is_autotune: Force re-run Triton autotuner and overwrite cache. Default False uses cached config.
+    :param skip_checks: Whether to skip input validation checks for faster performance.
+    :param return_lse: Whether to return the logsumexp tensor for numerical stability analysis. If True, returns a tuple (out, lse). If False, returns only out.
+    :returns: If return_lse is False, returns out with shape [batch_size, num_heads, head_dim]. If return_lse is True, returns a tuple (out, lse), where lse has shape [batch_size, num_heads].
+    """
+    if is_quant and (query_scale is None or key_scale is None or value_scale is None):
+        query, query_scale = quant.quantize_fp8(query)
+        key, key_scale = quant.quantize_fp8(key)
+        value, value_scale = quant.quantize_fp8(value)
+    if window_sizes is not None:
+        is_local = True
+    out, lse = _flash_sparse_attn_decode(
+        query=query,
+        key=key,
+        value=value,
+        softmax_scale=softmax_scale,
+        softmax_threshold=softmax_threshold,
+        query_scale=query_scale,
+        key_scale=key_scale,
+        value_scale=value_scale,
+        window_sizes=window_sizes,
+        is_local=is_local,
+        is_quant=is_quant,
+        out=out,
+        lse=lse,
+        is_autotune=is_autotune,
+        skip_checks=skip_checks,
+    )
+    if return_lse:
+        return out, lse
+    return out
+def flash_sparse_attn_varlen_func(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    is_causal: bool = False,
+    softmax_scale: Optional[float] = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    softmax_threshold: Optional[float] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    is_split_kv: bool = False,
+    is_split_qo: bool = False,
+    pack_gqa: bool = False,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+    return_lse: bool = False,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Flash sparse attention function for variable-length sequences that computes the attention output and optionally the logsumexp.
+    :param query: Query tensor of shape [total_seqlen_q, num_heads_q, head_dim].
+    :param key: Key tensor of shape [total_seqlen_k, num_heads_kv, head_dim].
+    :param value: Value tensor of shape [total_seqlen_k, num_heads_kv, head_dim].
+    :param cu_seqlens_q: Cumulative sequence lengths for queries, shape [batch_size + 1].
+    :param cu_seqlens_k: Cumulative sequence lengths for keys/values, shape [batch_size + 1].
+    :param max_seqlen_q: Maximum sequence length for queries.
+    :param max_seqlen_k: Maximum sequence length for keys/values.
+    :param is_causal: Whether to apply a causal mask.
+    :param softmax_scale: Optional scaling factor for the softmax. If None, defaults to 1/sqrt(head_dim).
+    :param query_scale: Optional per-tensor scale for query dequantization.
+    :param key_scale: Optional per-tensor scale for key dequantization.
+    :param value_scale: Optional per-tensor scale for value dequantization.
+    :param window_sizes: Optional window sizes tensor for flexible local attention. Must be shape [num_kv_heads, 4] with dtype int32, with columns [window_sink, window_left, window_right, window_dist]. If provided, is_local is automatically set to True.
+    :param softmax_threshold: Optional threshold for the sparse softmax. If None, defaults to head_dim / max_seqlen_k.
+    :param is_local: Whether to apply a local mask.
+    :param is_quant: Whether the inputs are quantized. If True, query_scale, key_scale, and value_scale must be provided or will be computed from the input tensors.
+    :param is_split_kv: Whether to enable split-KV for forward occupancy.
+    :param is_split_qo: Whether to enable split-QO for backward occupancy.
+    :param pack_gqa: Whether to pack grouped-query attention.
+    :param seqused_q: Optional tensor of shape [total_seqlen_q] indicating the actual sequence lengths for queries. If provided, overrides cu_seqlens_q for masking.
+    :param seqused_k: Optional tensor of shape [total_seqlen_k] indicating the actual sequence lengths for keys/values. If provided, overrides cu_seqlens_k for masking.
+    :param out: Optional preallocated output tensor with shape [batch_size, seqlen_q, num_heads, head_dim].
+    :param lse: Optional preallocated logsumexp tensor with shape [batch_size, num_heads, seqlen_q].
+    :param is_autotune: Force re-run Triton autotuner and overwrite cache. Default False uses cached config.
+    :param skip_checks: Whether to skip input validation checks for faster performance.
+    :param return_lse: Whether to return the logsumexp tensor for numerical stability analysis. If True, returns a tuple (out, lse). If False, returns only out.
+    :returns: If return_lse is False, returns out with shape [total_seqlen_q, num_heads_q, head_dim]. If return_lse is True, returns a tuple (out, lse), where lse has shape [total_seqlen_q, num_heads_q].
+    """
+    return FlashSparseAttnVarlenFunc.apply(
+        query,
+        key,
+        value,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        is_causal,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        softmax_threshold,
+        is_local,
+        is_quant,
+        is_split_kv,
+        is_split_qo,
+        pack_gqa,
+        seqused_q,
+        seqused_k,
+        out,
+        lse,
+        is_autotune,
+        skip_checks,
+        return_lse,
+    )
+def flash_sparse_attn_varlen_with_kvcache_func(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_k: int,
+    softmax_scale: Optional[float] = None,
+    softmax_threshold: Optional[float] = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    seqused_k: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+    return_lse: bool = False,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Flash sparse attention function for variable-length decoding with KV cache that computes the attention output and optionally the logsumexp.
+    :param query: Query tensor of shape [batch_size, num_heads_q, head_dim].
+    :param key: Key tensor of shape [total_seqlen_k, num_heads_kv, head_dim].
+    :param value: Value tensor of shape [total_seqlen_k, num_heads_kv, head_dim].
+    :param cu_seqlens_k: Cumulative sequence lengths for keys/values, shape [batch_size + 1].
+    :param max_seqlen_k: Maximum sequence length for keys/values.
+    :param softmax_scale: Optional scaling factor for the softmax. If None, defaults to 1/sqrt(head_dim).
+    :param softmax_threshold: Optional threshold for the sparse softmax. If None, defaults to head_dim / max_seqlen_k.
+    :param query_scale: Optional per-tensor scale for query dequantization.
+    :param key_scale: Optional per-tensor scale for key dequantization.
+    :param value_scale: Optional per-tensor scale for value dequantization.
+    :param window_sizes: Optional window sizes tensor for flexible local attention. Must be shape [num_kv_heads, 4] with dtype int32, with columns [window_sink, window_left, window_right, window_dist]. If provided, is_local is automatically set to True.
+    :param is_local: Whether to apply a local mask.
+    :param is_quant: Whether the inputs are quantized. If True, query_scale, key_scale, and value_scale must be provided for dequantization.
+    :param seqused_k: Optional tensor indicating the actual sequence lengths for keys/values.
+    :param out: Optional preallocated output tensor with shape [batch_size, num_heads_q, head_dim].
+    :param lse: Optional preallocated logsumexp tensor with shape [batch_size, num_heads_q].
+    :param is_autotune: Force re-run Triton autotuner and overwrite cache. Default False uses cached config.
+    :param skip_checks: Whether to skip input validation checks for faster performance.
+    :param return_lse: Whether to return the logsumexp tensor for numerical stability analysis. If True, returns a tuple (out, lse). If False, returns only out.
+    :returns: If return_lse is False, returns out with shape [batch_size, num_heads_q, head_dim]. If return_lse is True, returns a tuple (out, lse), where lse has shape [batch_size, num_heads_q].
+    """
+    if is_quant and (query_scale is None or key_scale is None or value_scale is None):
+        query, query_scale = quant.quantize_fp8(query)
+        key, key_scale = quant.quantize_fp8(key)
+        value, value_scale = quant.quantize_fp8(value)
+    if window_sizes is not None:
+        is_local = True
+    out, lse = _flash_sparse_attn_varlen_decode(
+        query=query,
+        key=key,
+        value=value,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_k=max_seqlen_k,
+        softmax_scale=softmax_scale,
+        softmax_threshold=softmax_threshold,
+        query_scale=query_scale,
+        key_scale=key_scale,
+        value_scale=value_scale,
+        window_sizes=window_sizes,
+        is_local=is_local,
+        is_quant=is_quant,
+        seqused_k=seqused_k,
+        out=out,
+        lse=lse,
+        is_autotune=is_autotune,
+        skip_checks=skip_checks,
+    )
+    if return_lse:
+        return out, lse
+    return out
+def flash_gated_attn_func(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    alpha: torch.Tensor,
+    delta: torch.Tensor,
+    is_causal: bool = False,
+    softmax_scale: Optional[float] = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    softmax_threshold: Optional[float] = None,
+    gate_threshold: Optional[float] = None,
+    is_logsigmoid_gate: bool = True,
+    is_adapt_gate: bool = True,
+    is_local: bool = False,
+    is_quant: bool = False,
+    is_split_kv: bool = False,
+    is_split_qo: bool = False,
+    pack_gqa: bool = False,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+    return_lse: bool = False,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Flash gated attention function that computes the attention output and optionally the logsumexp.
+    :param query: Query tensor of shape [batch_size, seqlen_q, num_heads, head_dim].
+    :param key: Key tensor of shape [batch_size, seqlen_k, num_kv_heads, head_dim].
+    :param value: Value tensor of shape [batch_size, seqlen_k, num_kv_heads, head_dim].
+    :param alpha: Tensor of shape [batch_size, seqlen_q, num_heads] representing the sparsity pattern for queries.
+    :param delta: Tensor of shape [batch_size, seqlen_k, num_kv_heads] representing the sparsity pattern for keys/values.
+    :param is_causal: Whether to apply a causal mask.
+    :param softmax_scale: Optional scaling factor for the softmax. If None, defaults to 1/sqrt(head_dim).
+    :param query_scale: Optional per-tensor scale for query dequantization.
+    :param key_scale: Optional per-tensor scale for key dequantization.
+    :param value_scale: Optional per-tensor scale for value dequantization.
+    :param window_sizes: Optional window sizes tensor for flexible local attention. Must be shape [num_kv_heads, 4] with dtype int32, with columns [window_sink, window_left, window_right, window_dist]. If provided, is_local is automatically set to True.
+    :param softmax_threshold: Optional threshold for the sparse softmax.
+    :param gate_threshold: Optional threshold for the sparsity gate.
+    :param is_logsigmoid_gate: Whether to use a log-sigmoid function for the sparsity gate. If False, uses a linear function.
+    :param is_adapt_gate: Whether to adapt the gate threshold based on sequence length.
+    :param is_local: Whether to apply a local mask.
+    :param is_quant: Whether the inputs are quantized. If True, query_scale, key_scale, and value_scale must be provided or will be computed from the input tensors.
+    :param is_split_kv: Whether to enable split-KV for forward occupancy.
+    :param is_split_qo: Whether to enable split-QO for backward occupancy.
+    :param pack_gqa: Whether to pack grouped-query attention.
+    :param out: Optional preallocated output tensor with shape [batch_size, seqlen_q, num_heads, head_dim].
+    :param lse: Optional preallocated logsumexp tensor with shape [batch_size, num_heads, seqlen_q].
+    :param is_autotune: Force re-run Triton autotuner and overwrite cache. Default False uses cached config.
+    :param skip_checks: Whether to skip input validation checks for faster performance.
+    :param return_lse: Whether to return the logsumexp tensor for numerical stability analysis. If True, returns a tuple (out, lse). If False, returns only out.
+    :returns: If return_lse is False, returns out with shape [batch_size, seqlen_q, num_heads, head_dim]. If return_lse is True, returns a tuple (out, lse), where lse has shape [batch_size, num_heads, seqlen_q].
+    """
+    return FlashGatedAttnFunc.apply(
+        query,
+        key,
+        value,
+        alpha,
+        delta,
+        is_causal,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        softmax_threshold,
+        gate_threshold,
+        is_logsigmoid_gate,
+        is_adapt_gate,
+        is_local,
+        is_quant,
+        is_split_kv,
+        is_split_qo,
+        pack_gqa,
+        out,
+        lse,
+        is_autotune,
+        skip_checks,
+        return_lse,
+    )
+def flash_gated_attn_with_kvcache_func(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    alpha: torch.Tensor,
+    delta: torch.Tensor,
+    softmax_scale: Optional[float] = None,
+    softmax_threshold: Optional[float] = None,
+    gate_threshold: Optional[float] = None,
+    is_logsigmoid_gate: bool = True,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+    return_lse: bool = False,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Flash gated attention function for decoding with KV cache that computes the attention output and optionally the logsumexp.
+    :param query: Query tensor of shape [batch_size, num_heads, head_dim].
+    :param key: Key tensor of shape [batch_size, seqlen_k, num_kv_heads, head_dim].
+    :param value: Value tensor of shape [batch_size, seqlen_k, num_kv_heads, head_dim].
+    :param alpha: Tensor of shape [batch_size, num_heads] representing the sparsity pattern for queries.
+    :param delta: Tensor of shape [batch_size, seqlen_k, num_kv_heads] representing the sparsity pattern for keys/values.
+    :param softmax_scale: Optional scaling factor for the softmax. If None, defaults to 1/sqrt(head_dim).
+    :param softmax_threshold: Optional threshold for the sparse softmax.
+    :param gate_threshold: Optional threshold for the sparsity gate.
+    :param is_logsigmoid_gate: Whether to use a log-sigmoid function for the sparsity gate. If False, uses a linear function.
+    :param query_scale: Optional per-tensor scale for query dequantization.
+    :param key_scale: Optional per-tensor scale for key dequantization.
+    :param value_scale: Optional per-tensor scale for value dequantization.
+    :param window_sizes: Optional window sizes tensor for flexible local attention. Must be shape [num_kv_heads, 4] with dtype int32, with columns [window_sink, window_left, window_right, window_dist]. If provided, is_local is automatically set to True.
+    :param is_local: Whether to apply a local mask.
+    :param is_quant: Whether the inputs are quantized. If True, query_scale, key_scale, and value_scale must be provided for dequantization.
+    :param out: Optional preallocated output tensor with shape [batch_size, num_heads, head_dim].
+    :param lse: Optional preallocated logsumexp tensor with shape [batch_size, num_heads].
+    :param is_autotune: Force re-run Triton autotuner and overwrite cache. Default False uses cached config.
+    :param skip_checks: Whether to skip input validation checks for faster performance.
+    :param return_lse: Whether to return the logsumexp tensor for numerical stability analysis. If True, returns a tuple (out, lse). If False, returns only out.
+    :returns: If return_lse is False, returns out with shape [batch_size, num_heads, head_dim]. If return_lse is True, returns a tuple (out, lse), where lse has shape [batch_size, num_heads].
+    """
+    if is_quant and (query_scale is None or key_scale is None or value_scale is None):
+        query, query_scale = quant.quantize_fp8(query)
+        key, key_scale = quant.quantize_fp8(key)
+        value, value_scale = quant.quantize_fp8(value)
+    if window_sizes is not None:
+        is_local = True
+    # Transpose for make_tensor_descriptor
+    delta = delta.transpose(-2, -1).contiguous()
+    out, lse = _flash_gated_attn_decode(
+        query=query,
+        key=key,
+        value=value,
+        alpha=alpha,
+        delta=delta,
+        softmax_scale=softmax_scale,
+        softmax_threshold=softmax_threshold,
+        gate_threshold=gate_threshold,
+        is_logsigmoid_gate=is_logsigmoid_gate,
+        query_scale=query_scale,
+        key_scale=key_scale,
+        value_scale=value_scale,
+        window_sizes=window_sizes,
+        is_local=is_local,
+        is_quant=is_quant,
+        out=out,
+        lse=lse,
+        is_autotune=is_autotune,
+        skip_checks=skip_checks,
+    )
+    if return_lse:
+        return out, lse
+    return out
+def flash_gated_attn_varlen_func(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    alpha: torch.Tensor,
+    delta: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    is_causal: bool = False,
+    softmax_scale: Optional[float] = None,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    softmax_threshold: Optional[float] = None,
+    gate_threshold: Optional[float] = None,
+    is_logsigmoid_gate: bool = True,
+    is_adapt_gate: bool = True,
+    is_local: bool = False,
+    is_quant: bool = False,
+    is_split_kv: bool = False,
+    is_split_qo: bool = False,
+    pack_gqa: bool = False,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+    return_lse: bool = False,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Flash gated attention function for variable-length sequences that computes the attention output and optionally the logsumexp.
+    :param query: Query tensor of shape [total_seqlen_q, num_heads_q, head_dim].
+    :param key: Key tensor of shape [total_seqlen_k, num_heads_kv, head_dim].
+    :param value: Value tensor of shape [total_seqlen_k, num_heads_kv, head_dim].
+    :param alpha: Tensor of shape [total_seqlen_q, num_heads_q] representing the sparsity pattern for queries.
+    :param delta: Tensor of shape [total_seqlen_k, num_heads_kv] representing the sparsity pattern for keys/values.
+    :param cu_seqlens_q: Cumulative sequence lengths for queries, shape [batch_size + 1].
+    :param cu_seqlens_k: Cumulative sequence lengths for keys/values, shape [batch_size + 1].
+    :param max_seqlen_q: Maximum sequence length for queries.
+    :param max_seqlen_k: Maximum sequence length for keys/values.
+    :param is_causal: Whether to apply a causal mask.
+    :param softmax_scale: Optional scaling factor for the softmax. If None, defaults to 1/sqrt(head_dim).
+    :param query_scale: Optional per-tensor scale for query dequantization.
+    :param key_scale: Optional per-tensor scale for key dequantization.
+    :param value_scale: Optional per-tensor scale for value dequantization.
+    :param window_sizes: Optional window sizes tensor for flexible local attention. Must be shape [num_kv_heads, 4] with dtype int32, with columns [window_sink, window_left, window_right, window_dist]. If provided, is_local is automatically set to True.
+    :param softmax_threshold: Optional threshold for the sparse softmax.
+    :param gate_threshold: Optional threshold for the sparsity gate.
+    :param is_logsigmoid_gate: Whether to use a log-sigmoid function for the sparsity gate. If False, uses a linear function.
+    :param is_adapt_gate: Whether to adapt the gate threshold based on sequence length.
+    :param is_local: Whether to apply a local mask.
+    :param is_quant: Whether the inputs are quantized. If True, query_scale, key_scale, and value_scale must be provided or will be computed from the input tensors.
+    :param is_split_kv: Whether to enable split-KV for forward occupancy.
+    :param is_split_qo: Whether to enable split-QO for backward occupancy.
+    :param pack_gqa: Whether to pack grouped-query attention.
+    :param seqused_q: Optional tensor of shape [total_seqlen_q] indicating the actual sequence lengths for queries. If provided, overrides cu_seqlens_q for masking.
+    :param seqused_k: Optional tensor of shape [total_seqlen_k] indicating the actual sequence lengths for keys/values. If provided, overrides cu_seqlens_k for masking.
+    :param out: Optional preallocated output tensor with shape [batch_size, seqlen_q, num_heads, head_dim].
+    :param lse: Optional preallocated logsumexp tensor with shape [batch_size, num_heads, seqlen_q].
+    :param is_autotune: Force re-run Triton autotuner and overwrite cache. Default False uses cached config.
+    :param skip_checks: Whether to skip input validation checks for faster performance.
+    :param return_lse: Whether to return the logsumexp tensor for numerical stability analysis. If True, returns a tuple (out, lse). If False, returns only out.
+    :returns: If return_lse is False, returns out with shape [total_seqlen_q, num_heads_q, head_dim]. If return_lse is True, returns a tuple (out, lse), where lse has shape [total_seqlen_q, num_heads_q].
+    """
+    return FlashGatedAttnVarlenFunc.apply(
+        query,
+        key,
+        value,
+        alpha,
+        delta,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        is_causal,
+        softmax_scale,
+        query_scale,
+        key_scale,
+        value_scale,
+        window_sizes,
+        softmax_threshold,
+        gate_threshold,
+        is_logsigmoid_gate,
+        is_adapt_gate,
+        is_local,
+        is_quant,
+        is_split_kv,
+        is_split_qo,
+        pack_gqa,
+        seqused_q,
+        seqused_k,
+        out,
+        lse,
+        is_autotune,
+        skip_checks,
+        return_lse,
+    )
+def flash_gated_attn_varlen_with_kvcache_func(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    alpha: torch.Tensor,
+    delta: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_k: int,
+    softmax_scale: Optional[float] = None,
+    softmax_threshold: Optional[float] = None,
+    gate_threshold: Optional[float] = None,
+    is_logsigmoid_gate: bool = True,
+    query_scale: Optional[torch.Tensor] = None,
+    key_scale: Optional[torch.Tensor] = None,
+    value_scale: Optional[torch.Tensor] = None,
+    window_sizes: Optional[torch.Tensor] = None,
+    is_local: bool = False,
+    is_quant: bool = False,
+    seqused_k: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+    lse: Optional[torch.Tensor] = None,
+    is_autotune: bool = False,
+    skip_checks: bool = False,
+    return_lse: bool = False,
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Flash gated attention function for variable-length decoding with KV cache that computes the attention output and optionally the logsumexp.
+    :param query: Query tensor of shape [batch_size, num_heads_q, head_dim].
+    :param key: Key tensor of shape [total_seqlen_k, num_heads_kv, head_dim].
+    :param value: Value tensor of shape [total_seqlen_k, num_heads_kv, head_dim].
+    :param alpha: Tensor of shape [batch_size, num_heads_q] representing the sparsity pattern for queries.
+    :param delta: Tensor of shape [total_seqlen_k, num_heads_kv] representing the sparsity pattern for keys/values.
+    :param cu_seqlens_k: Cumulative sequence lengths for keys/values, shape [batch_size + 1].
+    :param max_seqlen_k: Maximum sequence length for keys/values.
+    :param softmax_scale: Optional scaling factor for the softmax. If None, defaults to 1/sqrt(head_dim).
+    :param softmax_threshold: Optional threshold for the sparse softmax.
+    :param gate_threshold: Optional threshold for the sparsity gate.
+    :param is_logsigmoid_gate: Whether to use a log-sigmoid function for the sparsity gate. If False, uses a linear function.
+    :param query_scale: Optional per-tensor scale for query dequantization.
+    :param key_scale: Optional per-tensor scale for key dequantization.
+    :param value_scale: Optional per-tensor scale for value dequantization.
+    :param window_sizes: Optional window sizes tensor for flexible local attention. Must be shape [num_kv_heads, 4] with dtype int32, with columns [window_sink, window_left, window_right, window_dist]. If provided, is_local is automatically set to True.
+    :param is_local: Whether to apply a local mask.
+    :param is_quant: Whether the inputs are quantized. If True, query_scale, key_scale, and value_scale must be provided for dequantization.
+    :param seqused_k: Optional tensor indicating the actual sequence lengths for keys/values.
+    :param out: Optional preallocated output tensor with shape [batch_size, num_heads_q, head_dim].
+    :param lse: Optional preallocated logsumexp tensor with shape [batch_size, num_heads_q].
+    :param is_autotune: Force re-run Triton autotuner and overwrite cache. Default False uses cached config.
+    :param skip_checks: Whether to skip input validation checks for faster performance.
+    :param return_lse: Whether to return the logsumexp tensor for numerical stability analysis. If True, returns a tuple (out, lse). If False, returns only out.
+    :returns: If return_lse is False, returns out with shape [batch_size, num_heads_q, head_dim]. If return_lse is True, returns a tuple (out, lse), where lse has shape [batch_size, num_heads_q].
+    """
+    if is_quant and (query_scale is None or key_scale is None or value_scale is None):
+        query, query_scale = quant.quantize_fp8(query)
+        key, key_scale = quant.quantize_fp8(key)
+        value, value_scale = quant.quantize_fp8(value)
+    if window_sizes is not None:
+        is_local = True
+    # Transpose for make_tensor_descriptor
+    delta = delta.transpose(-2, -1).contiguous()
+    out, lse = _flash_gated_attn_varlen_decode(
+        query=query,
+        key=key,
+        value=value,
+        alpha=alpha,
+        delta=delta,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_k=max_seqlen_k,
+        softmax_scale=softmax_scale,
+        softmax_threshold=softmax_threshold,
+        gate_threshold=gate_threshold,
+        is_logsigmoid_gate=is_logsigmoid_gate,
+        query_scale=query_scale,
+        key_scale=key_scale,
+        value_scale=value_scale,
+        window_sizes=window_sizes,
+        is_local=is_local,
+        is_quant=is_quant,
+        seqused_k=seqused_k,
+        out=out,
+        lse=lse,
+        is_autotune=is_autotune,
+        skip_checks=skip_checks,
+    )
+    if return_lse:
+        return out, lse
+    return out
+def flash_proj_func(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    out_dtype: torch.dtype = torch.float8_e5m2,
+    prev_scale: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Fused linear projection with quantized output and per-tensor scale.
+    :param input: Input tensor of shape [batch_size, seqlen, in_features].
+    :param weight: Weights tensor of shape [out_features, in_features].
+    :param out_dtype: Target output dtype.
+    :param prev_scale: Per-tensor scale from the previous forward.
+    :returns: Tuple of (out, actual_scale), where out has shape [batch_size, seqlen, out_features].
+    """
+    return FlashProjFunc.apply(input, weight, out_dtype, prev_scale)

build/torch-cuda/kernel_repr.py ADDED Viewed

	@@ -0,0 +1,115 @@

+def _get(c, key, default=None):
+    try:
+        return c[key]
+    except (KeyError, TypeError):
+        return default
+def fwd_dense_repr(specialization):
+    c = specialization.constants
+    mask = (
+        "causal"
+        if _get(c, "IS_CAUSAL")
+        else ("local" if _get(c, "IS_LOCAL") else "full")
+    )
+    varlen = "_varlen" if _get(c, "HAS_CU_SEQLENS_Q") else ""
+    split = "_splitkv" if _get(c, "IS_SPLIT_KV") else ""
+    return f"flash_dense_fwd_{mask}_{c['TILE_M']}x{c['TILE_N']}{split}{varlen}"
+def fwd_sparse_repr(specialization):
+    c = specialization.constants
+    mask = (
+        "causal"
+        if _get(c, "IS_CAUSAL")
+        else ("local" if _get(c, "IS_LOCAL") else "full")
+    )
+    varlen = "_varlen" if _get(c, "HAS_CU_SEQLENS_Q") else ""
+    split = "_splitkv" if _get(c, "IS_SPLIT_KV") else ""
+    return f"flash_sparse_fwd_{mask}_{c['TILE_M']}x{c['TILE_N']}{split}{varlen}"
+def fwd_gated_repr(specialization):
+    c = specialization.constants
+    mask = (
+        "causal"
+        if _get(c, "IS_CAUSAL")
+        else ("local" if _get(c, "IS_LOCAL") else "full")
+    )
+    varlen = "_varlen" if _get(c, "HAS_CU_SEQLENS_Q") else ""
+    split = "_splitkv" if _get(c, "IS_SPLIT_KV") else ""
+    return f"flash_gated_fwd_{mask}_{c['TILE_M']}x{c['TILE_N']}{split}{varlen}"
+def bwd_dense_repr(specialization):
+    c = specialization.constants
+    mask = (
+        "causal"
+        if _get(c, "IS_CAUSAL")
+        else ("local" if _get(c, "IS_LOCAL") else "full")
+    )
+    varlen = "_varlen" if _get(c, "HAS_CU_SEQLENS_Q") else ""
+    return f"flash_dense_bwd_{mask}_{c['TILE_M']}x{c['TILE_N']}{varlen}"
+def bwd_sparse_repr(specialization):
+    c = specialization.constants
+    mask = (
+        "causal"
+        if _get(c, "IS_CAUSAL")
+        else ("local" if _get(c, "IS_LOCAL") else "full")
+    )
+    varlen = "_varlen" if _get(c, "HAS_CU_SEQLENS_Q") else ""
+    return f"flash_sparse_bwd_{mask}_{c['TILE_M']}x{c['TILE_N']}{varlen}"
+def bwd_gated_repr(specialization):
+    c = specialization.constants
+    mask = (
+        "causal"
+        if _get(c, "IS_CAUSAL")
+        else ("local" if _get(c, "IS_LOCAL") else "full")
+    )
+    varlen = "_varlen" if _get(c, "HAS_CU_SEQLENS_Q") else ""
+    return f"flash_gated_bwd_{mask}_{c['TILE_M']}x{c['TILE_N']}{varlen}"
+def dec_dense_repr(specialization):
+    c = specialization.constants
+    mask = "local" if _get(c, "IS_LOCAL") else "full"
+    varlen = "_varlen" if _get(c, "HAS_CU_SEQLENS_Q") else ""
+    return f"flash_dense_dec_{mask}_{c['TILE_M']}x{c['TILE_N']}{varlen}"
+def dec_sparse_repr(specialization):
+    c = specialization.constants
+    mask = "local" if _get(c, "IS_LOCAL") else "full"
+    varlen = "_varlen" if _get(c, "HAS_CU_SEQLENS_Q") else ""
+    return f"flash_sparse_dec_{mask}_{c['TILE_M']}x{c['TILE_N']}{varlen}"
+def dec_gated_repr(specialization):
+    c = specialization.constants
+    mask = "local" if _get(c, "IS_LOCAL") else "full"
+    varlen = "_varlen" if _get(c, "HAS_CU_SEQLENS_Q") else ""
+    return f"flash_gated_dec_{mask}_{c['TILE_M']}x{c['TILE_N']}{varlen}"
+def fwd_combine_repr(specialization):
+    c = specialization.constants
+    return f"flash_fwd_combine_{c['TILE_M']}"
+def dec_combine_repr(specialization):
+    c = specialization.constants
+    return f"flash_dec_combine_{c['TILE_K']}"
+def bwd_preprocess_repr(specialization):
+    c = specialization.constants
+    return f"flash_bwd_preprocess_{c['TILE_M']}"
+def bwd_postprocess_repr(specialization):
+    c = specialization.constants
+    return f"flash_bwd_postprocess_{c['TILE_M']}"

build/torch-cuda/launch_grid.py ADDED Viewed

	@@ -0,0 +1,202 @@

+import triton
+from . import cache_utils
+def get_fwd_grid(
+    batch_size: int,
+    seqlen_q: int,
+    num_heads_q: int,
+    num_heads_kv: int,
+    pack_gqa: bool,
+    num_splits: int,
+):
+    """
+    Get the grid function for the forward kernel.
+    :param batch_size: Batch size
+    :param seqlen_q: Sequence length of queries
+    :param num_heads_q: Number of query heads
+    :param num_heads_kv: Number of key/value heads
+    :param pack_gqa: Whether GQA packing is used
+    :param num_splits: Number of KV splits
+    :return grid: Grid function
+    """
+    def grid(META):
+        return (
+            triton.cdiv(
+                seqlen_q * (num_heads_q // num_heads_kv) if pack_gqa else seqlen_q,
+                META["TILE_M"],
+            ),
+            num_heads_kv if pack_gqa else num_heads_q,
+            batch_size * num_splits,
+        )
+    return grid
+get_fwd_grid = cache_utils.cache_launch_grid(get_fwd_grid)
+def get_bwd_grid(
+    batch_size: int,
+    seqlen_k: int,
+    num_heads_q: int,
+    num_splits: int = 1,
+):
+    """
+    Get the grid function for the backward kernel.
+    :param batch_size: Batch size
+    :param seqlen_k: Sequence length of keys
+    :param num_heads_q: Number of query heads
+    :param num_splits: Number of QO splits
+    :return grid: Grid function
+    """
+    def grid(META):
+        return (
+            triton.cdiv(seqlen_k, META["TILE_N"]),
+            num_heads_q,
+            batch_size * num_splits,
+        )
+    return grid
+get_bwd_grid = cache_utils.cache_launch_grid(get_bwd_grid)
+def get_dec_grid(
+    batch_size: int,
+    num_heads_kv: int,
+    num_splits: int,
+):
+    """
+    Get the grid function for the decode kernel.
+    :param batch_size: Batch size
+    :param num_heads_kv: Number of key/value heads
+    :param num_splits: Number of KV splits
+    :return grid: Grid function
+    """
+    def grid(META):
+        return (
+            num_heads_kv,
+            batch_size * num_splits,
+        )
+    return grid
+get_dec_grid = cache_utils.cache_launch_grid(get_dec_grid)
+def get_fwd_combine_grid(
+    batch_size: int,
+    seqlen_q: int,
+    num_heads_q: int,
+):
+    """
+    Get the grid function for the forward combine kernel.
+    :param batch_size: Batch size
+    :param seqlen_q: Sequence length of queries
+    :param num_heads_q: Number of query heads
+    :return grid: Grid function
+    """
+    def grid(META):
+        return (
+            triton.cdiv(seqlen_q, META["TILE_M"]),
+            batch_size * num_heads_q,
+        )
+    return grid
+get_fwd_combine_grid = cache_utils.cache_launch_grid(get_fwd_combine_grid)
+def get_bwd_preprocess_grid(
+    batch_size: int,
+    seqlen_q: int,
+    num_heads_q: int,
+):
+    """
+    Get the grid function for the backward preprocess kernel.
+    :param batch_size: Batch size
+    :param seqlen_q: Sequence length of queries
+    :param num_heads_q: Number of query heads
+    :return grid: Grid function
+    """
+    def grid(META):
+        return (
+            triton.cdiv(seqlen_q, META["TILE_M"]),
+            num_heads_q,
+            batch_size,
+        )
+    return grid
+get_bwd_preprocess_grid = cache_utils.cache_launch_grid(get_bwd_preprocess_grid)
+def get_bwd_postprocess_grid(
+    batch_size: int,
+    seqlen_q: int,
+    num_heads_q: int,
+):
+    """
+    Get the grid function for the backward postprocess kernel.
+    :param batch_size: Batch size
+    :param seqlen_q: Sequence length of queries
+    :param num_heads_q: Number of query heads
+    :return grid: Grid function
+    """
+    def grid(META):
+        return (
+            triton.cdiv(seqlen_q, META["TILE_M"]),
+            num_heads_q,
+            batch_size,
+        )
+    return grid
+get_bwd_postprocess_grid = cache_utils.cache_launch_grid(get_bwd_postprocess_grid)
+def get_dec_combine_grid(
+    batch_size: int,
+    num_heads_q: int,
+):
+    """
+    Get the grid function for the decode combine kernel.
+    :param batch_size: Batch size
+    :param num_heads_q: Number of query heads
+    :return grid: Grid function
+    """
+    def grid(META):
+        return (batch_size * num_heads_q,)
+    return grid
+get_dec_combine_grid = cache_utils.cache_launch_grid(get_dec_combine_grid)

build/torch-cuda/launch_template.py ADDED Viewed

	@@ -0,0 +1,457 @@

+"""
+NOTE: Setting num_ctas=2 for the kernel can trigger Triton's PlanCTA assertion
+Setting num_ctas=1 for now to avoid this issue, but we may want to revisit this in the future
+"""
+import functools
+import hashlib
+import json
+import os
+import sys
+import tempfile
+from pathlib import Path
+from filelock import ReadWriteLock
+import torch
+import triton
+from . import cache_utils
+TRITON_CACHE_DIR: str = os.getenv(
+    "FLASH_SPARSE_ATTENTION_LAUNCH_CONFIG_DIR",
+    str(Path(__file__).resolve().parent / "launch_config"),
+)
+@functools.lru_cache(maxsize=1)
+def _compute_source_fingerprint() -> str:
+    """
+    Hash all Triton Python sources plus runtime ABI stamps into a short fingerprint.
+    The fingerprint changes whenever any .py file under flash_sparse_attn/ops/triton is
+    modified, the Python minor version changes, or the triton package version changes.
+    :return: SHA-256 hex digest string.
+    """
+    triton_root = Path(__file__).resolve().parent
+    h = hashlib.sha256()
+    h.update(f"py{sys.version_info.major}.{sys.version_info.minor}".encode())
+    h.update(f"triton={getattr(triton, '__version__', 'unknown')}".encode())
+    for src in sorted(triton_root.rglob("*.py")):
+        if not src.is_file():
+            continue
+        h.update(src.relative_to(triton_root).as_posix().encode())
+        content = src.read_bytes()
+        h.update(len(content).to_bytes(8, "little"))
+        h.update(content)
+    return h.hexdigest()
+def _seqlen_bucket(seqlen: int) -> int:
+    if seqlen <= 0:
+        return 1
+    return triton.next_power_of_2(seqlen)
+def _sanitize_device_name(name: str) -> str:
+    return name.replace(" ", "_").replace("/", "_")
+class LaunchConfigCache:
+    """
+    Persistent per-device cache for Triton kernel launch configurations.
+    Stores tuned (TILE_M, TILE_N, num_warps, num_stages, num_ctas) keyed by
+    kernel name, sequence length buckets, and dispatch flags. Cache is namespaced
+    under a source fingerprint directory for automatic invalidation on code changes.
+    """
+    __slots__ = ("_memory", "_fingerprint")
+    def __init__(self):
+        self._memory: dict[str, dict[tuple, tuple]] = {}
+        self._fingerprint: str = _compute_source_fingerprint()
+    def _get_cache_dir(self) -> Path:
+        cache_dir = Path(TRITON_CACHE_DIR) / self._fingerprint
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        return cache_dir
+    def _json_path(self, device_name: str) -> Path:
+        return self._get_cache_dir() / f"{_sanitize_device_name(device_name)}.json"
+    def _lock(self, device_name: str) -> ReadWriteLock:
+        lock_path = self._get_cache_dir() / f"{_sanitize_device_name(device_name)}"
+        return ReadWriteLock(str(lock_path), timeout=15)
+    def _read_unlocked(self, device_name: str) -> dict[tuple, tuple]:
+        json_path = self._json_path(device_name)
+        result: dict[tuple, tuple] = {}
+        if not json_path.exists():
+            return result
+        try:
+            data = json.loads(json_path.read_text())
+            for entry in data:
+                key = (
+                    entry["kernel"],
+                    entry["seqlen_q_bucket"],
+                    entry["seqlen_k_bucket"],
+                    entry["tile_k"],
+                    entry["is_local"],
+                    entry["qhead_per_kvhead"],
+                    entry.get("is_causal", False),
+                    entry.get("pack_gqa", False),
+                    entry.get("is_quant", False),
+                )
+                result[key] = tuple(entry["config"])
+        except (OSError, json.JSONDecodeError, KeyError, TypeError):
+            result = {}
+        return result
+    def _write_unlocked(self, device_name: str, cache: dict[tuple, tuple]) -> None:
+        json_path = self._json_path(device_name)
+        data = []
+        for (k_name, sq, sk, tk, local, qpk, causal, pgqa, quant), config in sorted(
+            cache.items(), key=lambda x: x[0]
+        ):
+            data.append(
+                {
+                    "kernel": k_name,
+                    "seqlen_q_bucket": sq,
+                    "seqlen_k_bucket": sk,
+                    "tile_k": tk,
+                    "is_local": local,
+                    "qhead_per_kvhead": qpk,
+                    "is_causal": causal,
+                    "pack_gqa": pgqa,
+                    "is_quant": quant,
+                    "config": list(config),
+                }
+            )
+        fd, tmp_path = tempfile.mkstemp(
+            dir=str(json_path.parent), suffix=".json.tmp", prefix=json_path.stem
+        )
+        try:
+            with os.fdopen(fd, "w") as f:
+                json.dump(data, f)
+            os.replace(tmp_path, str(json_path))
+        except BaseException:
+            try:
+                os.unlink(tmp_path)
+            except OSError:
+                pass
+            raise
+    def _load_device(self, device_name: str) -> dict[tuple, tuple]:
+        if device_name in self._memory:
+            return self._memory[device_name]
+        with self._lock(device_name).read_lock():
+            result = self._read_unlocked(device_name)
+        self._memory[device_name] = result
+        return result
+    def get(
+        self,
+        device: torch.device,
+        kernel_name: str,
+        seqlen_q: int,
+        seqlen_k: int,
+        tile_k: int,
+        is_local: bool = False,
+        qhead_per_kvhead: int = 1,
+        is_causal: bool = False,
+        pack_gqa: bool = False,
+        is_quant: bool = False,
+    ) -> tuple[int, int, int, int, int] | None:
+        """
+        Load a cached launch config for the given kernel specialization.
+        :param device: The device to run the kernel on
+        :param kernel_name: Name of the Triton kernel
+        :param seqlen_q: Sequence length of queries
+        :param seqlen_k: Sequence length of keys
+        :param tile_k: Tile size along the K dimension
+        :param is_local: Whether local mask is applied
+        :param qhead_per_kvhead: Ratio of query heads to key/value heads
+        :param is_causal: Whether causal mask is applied
+        :param pack_gqa: Whether GQA packing is enabled
+        :param is_quant: Whether quantization is used
+        :return launch_config: Tuple of (TILE_M, TILE_N, num_warps, num_stages, num_ctas) or None on cache miss
+        """
+        device_name = torch.cuda.get_device_name(device)
+        cache = self._load_device(device_name)
+        key = (
+            kernel_name,
+            _seqlen_bucket(seqlen_q),
+            _seqlen_bucket(seqlen_k),
+            tile_k,
+            is_local,
+            qhead_per_kvhead,
+            is_causal,
+            pack_gqa,
+            is_quant,
+        )
+        return cache.get(key)
+    def put(
+        self,
+        device: torch.device,
+        kernel_name: str,
+        seqlen_q: int,
+        seqlen_k: int,
+        tile_k: int,
+        is_local: bool = False,
+        qhead_per_kvhead: int = 1,
+        is_causal: bool = False,
+        pack_gqa: bool = False,
+        is_quant: bool = False,
+        *,
+        config: tuple[int, int, int, int, int],
+    ) -> None:
+        """
+        Store a launch config to the per-device cache.
+        :param device: The device to run the kernel on
+        :param kernel_name: Name of the Triton kernel
+        :param seqlen_q: Sequence length of queries
+        :param seqlen_k: Sequence length of keys
+        :param tile_k: Tile size along the K dimension
+        :param is_local: Whether local mask is applied
+        :param qhead_per_kvhead: Ratio of query heads to key/value heads
+        :param is_causal: Whether causal mask is applied
+        :param pack_gqa: Whether GQA packing is enabled
+        :param is_quant: Whether quantization is used
+        :param config: Tuple of (TILE_M, TILE_N, num_warps, num_stages, num_ctas)
+        """
+        device_name = torch.cuda.get_device_name(device)
+        key = (
+            kernel_name,
+            _seqlen_bucket(seqlen_q),
+            _seqlen_bucket(seqlen_k),
+            tile_k,
+            is_local,
+            qhead_per_kvhead,
+            is_causal,
+            pack_gqa,
+            is_quant,
+        )
+        # Skip disk write if memory cache already has the same config
+        mem = self._memory.get(device_name)
+        if mem is not None and mem.get(key) == config:
+            return
+        with self._lock(device_name).write_lock():
+            disk_cache = self._read_unlocked(device_name)
+            disk_cache[key] = config
+            self._write_unlocked(device_name, disk_cache)
+        self._memory[device_name] = disk_cache
+    def clear(self) -> None:
+        self._memory.clear()
+    def stats(self) -> dict:
+        total_entries = sum(len(v) for v in self._memory.values())
+        return {
+            "devices": len(self._memory),
+            "entries": total_entries,
+            "fingerprint": self._fingerprint,
+        }
+_cache_instance: LaunchConfigCache | None = None
+def get_launch_config_cache() -> LaunchConfigCache:
+    global _cache_instance
+    if _cache_instance is None:
+        _cache_instance = LaunchConfigCache()
+    return _cache_instance
+def load_launch_config(
+    device: torch.device,
+    kernel_name: str,
+    seqlen_q: int,
+    seqlen_k: int,
+    tile_k: int,
+    is_local: bool = False,
+    qhead_per_kvhead: int = 1,
+    is_causal: bool = False,
+    pack_gqa: bool = False,
+    is_quant: bool = False,
+) -> tuple[int, int, int, int, int] | None:
+    """
+    Load cached launch config for a kernel specialization.
+    :param device: The device to run the kernel on
+    :param kernel_name: Name of the Triton kernel
+    :param seqlen_q: Sequence length of queries
+    :param seqlen_k: Sequence length of keys
+    :param tile_k: Tile size along the K dimension
+    :param is_local: Whether local mask is applied
+    :param qhead_per_kvhead: Ratio of query heads to key/value heads
+    :param is_causal: Whether causal mask is applied
+    :param pack_gqa: Whether GQA packing is enabled
+    :param is_quant: Whether quantization is used
+    :return launch_config: Tuple of (TILE_M, TILE_N, num_warps, num_stages, num_ctas) or None on cache miss
+    """
+    return get_launch_config_cache().get(
+        device,
+        kernel_name,
+        seqlen_q,
+        seqlen_k,
+        tile_k,
+        is_local,
+        qhead_per_kvhead,
+        is_causal,
+        pack_gqa,
+        is_quant,
+    )
+def store_launch_config(
+    device: torch.device,
+    kernel_name: str,
+    seqlen_q: int,
+    seqlen_k: int,
+    tile_k: int,
+    *,
+    config: tuple[int, int, int, int, int],
+    is_local: bool = False,
+    qhead_per_kvhead: int = 1,
+    is_causal: bool = False,
+    pack_gqa: bool = False,
+    is_quant: bool = False,
+) -> None:
+    """
+    Store a launch config to the per-device JSON cache.
+    :param device: The device to run the kernel on
+    :param kernel_name: Name of the Triton kernel
+    :param seqlen_q: Sequence length of queries
+    :param seqlen_k: Sequence length of keys
+    :param tile_k: Tile size along the K dimension
+    :param config: Tuple of (TILE_M, TILE_N, num_warps, num_stages, num_ctas)
+    :param is_local: Whether local mask is applied
+    :param qhead_per_kvhead: Ratio of query heads to key/value heads
+    :param is_causal: Whether causal mask is applied
+    :param pack_gqa: Whether GQA packing is enabled
+    :param is_quant: Whether quantization is used
+    """
+    get_launch_config_cache().put(
+        device,
+        kernel_name,
+        seqlen_q,
+        seqlen_k,
+        tile_k,
+        is_local,
+        qhead_per_kvhead,
+        is_causal,
+        pack_gqa,
+        is_quant,
+        config=config,
+    )
+def extract_best_config(autotuned_kernel) -> tuple[int, int, int, int, int] | None:
+    """
+    Extract best config from a Triton autotuned kernel after it has run.
+    :param autotuned_kernel: A Triton autotuned kernel instance
+    :return launch_config: Tuple of (TILE_M, TILE_N, num_warps, num_stages, num_ctas) or None if no best config found
+    """
+    raw = getattr(autotuned_kernel, "_autotuned", autotuned_kernel)
+    best = getattr(raw, "best_config", None)
+    if best is None:
+        return None
+    return (
+        best.kwargs.get("TILE_M", 64),
+        best.kwargs.get("TILE_N", 64),
+        best.num_warps,
+        best.num_stages,
+        getattr(best, "num_ctas", 1),
+    )
+def get_fwd_combine_launch_config(
+    tile_k: int,
+    device: torch.device,
+    arch: int,
+) -> tuple[int, int, int, int]:
+    """
+    Get launch configuration for forward combine kernel based on input parameters and device architecture.
+    :param tile_k: Tile size along the K dimension
+    :param device: The device to run the kernel on
+    :param arch: The architecture of the device
+    :return launch_config: Tuple of (tile_m, num_warps, num_stages, num_ctas) for launching the kernel
+    """
+    if arch == -1:
+        raise NotImplementedError(f"Unsupported device: {device} with arch {arch}")
+    if device.type == "cuda":
+        if arch // 10 == 8:
+            tile_m = 4 if tile_k % 128 == 0 else (8 if tile_k % 64 == 0 else 16)
+            return (tile_m, 4, 1, 1)
+        elif arch // 10 == 9:
+            tile_m = 8 if tile_k % 128 == 0 else (16 if tile_k % 64 == 0 else 32)
+            return (tile_m, 4, 1, 1)
+        elif arch // 10 == 10:
+            tile_m = 16 if tile_k % 128 == 0 else (32 if tile_k % 64 == 0 else 64)
+            return (tile_m, 4, 1, 1)
+        elif arch // 10 == 12:
+            tile_m = 4 if tile_k % 128 == 0 else (8 if tile_k % 64 == 0 else 16)
+            return (tile_m, 4, 1, 1)
+        else:
+            raise NotImplementedError(f"Unsupported CUDA architecture: {arch}")
+    else:
+        raise NotImplementedError(f"Unsupported device type: {device.type}")
+get_fwd_combine_launch_config = cache_utils.cache_launch_config(
+    get_fwd_combine_launch_config
+)
+def get_dec_combine_launch_config(
+    tile_k: int,
+    device: torch.device,
+    arch: int,
+) -> tuple[int, int, int]:
+    """
+    Get launch configuration for decode combine kernel based on input parameters and device architecture.
+    :param tile_k: Tile size along the K dimension
+    :param device: The device to run the kernel on
+    :param arch: The architecture of the device
+    :return launch_config: Tuple of (num_warps, num_stages, num_ctas) for launching the kernel
+    """
+    if arch == -1:
+        raise NotImplementedError(f"Unsupported device: {device} with arch {arch}")
+    if device.type == "cuda":
+        if arch // 10 == 8:
+            num_stages = max(min(8, 512 // tile_k), 1)
+            return (4, num_stages, 1)
+        elif arch // 10 == 9:
+            num_stages = max(min(8, 512 // tile_k), 1)
+            return (4, num_stages, 1)
+        elif arch // 10 == 10:
+            num_stages = max(min(16, 512 // tile_k), 1)
+            return (4, num_stages, 1)
+        elif arch // 10 == 12:
+            num_stages = max(min(4, 512 // tile_k), 1)
+            return (4, num_stages, 1)
+        else:
+            raise NotImplementedError(f"Unsupported CUDA architecture: {arch}")
+    else:
+        raise NotImplementedError(f"Unsupported device type: {device.type}")
+get_dec_combine_launch_config = cache_utils.cache_launch_config(
+    get_dec_combine_launch_config
+)

build/torch-cuda/mask.py ADDED Viewed

	@@ -0,0 +1,106 @@

+import triton
+import triton.language as tl
+@triton.jit
+def apply_mask(
+    acc_s,
+    m_block,
+    n_block,
+    seqlen_q,
+    seqlen_k,
+    window_size_sink,
+    window_size_left,
+    window_size_right,
+    window_size_dist,
+    MASK_SEQLEN: tl.constexpr,
+    MASK_CAUSAL: tl.constexpr,
+    MASK_LOCAL: tl.constexpr,
+    MASK_SINK: tl.constexpr,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+    QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr,
+    SWAP_AB: tl.constexpr,
+):
+    """
+    Apply seqlen, causal, and local masks to the attention scores.
+    :param acc_s: Attention scores tensor of shape [BLOCK_M, BLOCK_N].
+    :param m_block: Current block index along the M dimension.
+    :param n_block: Current block index along the N dimension.
+    :param seqlen_q: The sequence length of the query.
+    :param seqlen_k: The sequence length of the key.
+    :param window_size_sink: Prefix sink token count.
+    :param window_size_left: Distant local band token count.
+    :param window_size_right: Gap token count after the near-diagonal window before the distant band.
+    :param window_size_dist: Near-diagonal local token count.
+    :param MASK_SEQLEN: Boolean flag indicating if seqlen masking should be applied.
+    :param MASK_CAUSAL: Boolean flag indicating if causal masking should be applied.
+    :param MASK_LOCAL: Boolean flag indicating if local masking should be applied.
+    :param MASK_SINK: Boolean flag indicating if sink masking should be applied.
+    :param TILE_M: Tile size along the M dimension.
+    :param TILE_N: Tile size along the N dimension.
+    :param QHEAD_PER_KVHEAD_PACKGQA: Ratio of query heads to key/value heads for packed GQA.
+    :param SWAP_AB: Boolean flag indicating if query and key dimensions are swapped.
+    :return acc_s: Masked attention scores tensor of shape [BLOCK_M, BLOCK_N].
+    """
+    offs_m = m_block * TILE_M + tl.arange(0, TILE_M)
+    offs_n = n_block * TILE_N + tl.arange(0, TILE_N)
+    if SWAP_AB:
+        tl.static_assert(
+            QHEAD_PER_KVHEAD_PACKGQA == 1, "SWAP_AB with PACKGQA > 1 not supported"
+        )
+        q_idx = offs_m[None, :]
+        k_idx = offs_n[:, None]
+    else:
+        q_idx = offs_m[:, None]
+        k_idx = offs_n[None, :]
+        if QHEAD_PER_KVHEAD_PACKGQA > 1:
+            q_idx = q_idx // QHEAD_PER_KVHEAD_PACKGQA
+    if MASK_SEQLEN:
+        acc_s = tl.where(
+            (k_idx < seqlen_k) & (q_idx < seqlen_q),
+            acc_s,
+            float("-inf"),
+        )
+    if MASK_CAUSAL or MASK_LOCAL or MASK_SINK:
+        causal_offset = seqlen_k - seqlen_q
+        dist = q_idx + causal_offset - k_idx
+        if MASK_LOCAL or MASK_SINK:
+            allowed = (dist >= 0) & (dist < 0)
+            if MASK_LOCAL:
+                allowed = allowed | ((dist >= 0) & (dist < window_size_dist))
+                allowed = allowed | (
+                    (dist >= window_size_dist + window_size_right)
+                    & (dist < window_size_dist + window_size_right + window_size_left)
+                )
+            if MASK_SINK:
+                allowed = allowed | ((k_idx < window_size_sink) & (dist >= 0))
+            acc_s = tl.where(
+                allowed,
+                acc_s,
+                float("-inf"),
+            )
+        elif MASK_CAUSAL:
+            acc_s = tl.where(
+                dist >= 0,
+                acc_s,
+                float("-inf"),
+            )
+        else:
+            acc_s = tl.where(
+                ((dist >= 0) & (dist < window_size_dist))
+                | (
+                    (dist >= window_size_dist + window_size_right)
+                    & (dist < window_size_dist + window_size_right + window_size_left)
+                ),
+                acc_s,
+                float("-inf"),
+            )
+    return acc_s

build/torch-cuda/metadata.json ADDED Viewed

	@@ -0,0 +1,47 @@

+{
+  "name": "flash-sparse-attention",
+  "id": "_flash_sparse_attention_cuda_7db628a_dirty",
+  "version": 1,
+  "license": "BSD-3-Clause",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "OXYzVpCpaU4x34L+qzPOSlycSDR+2i/CsF04RT/tBL4=",
+      "_ops.py": "Ds3MXWORK2pW+eBqjHYPr5gUiUF/1eJs3P2ryUEnuH4=",
+      "activations.py": "bAa6rHqI0noI/8zSGNrY2uzM5VHQmnsS4d2b+V0e5eU=",
+      "assert_inputs.py": "eLNN5h2Qhq3hdg2nnSSOAqLUqMZcYW2QqFvPo1FGmLg=",
+      "autotuner.py": "rpAY5s+/E6gSau/zZrNJjVOmvteEHsOjDNX/qAsc1Z0=",
+      "block_info.py": "JCm1Pp0jCHJOFGpwK6eXdiNhr2ZospuXx8bAX7Bez1E=",
+      "cache_utils.py": "+PzgjtFoef+lFNqrVvst06iOLmkTOXeW/y95kFOEi64=",
+      "flash_bwd_postprocess.py": "HqYiH4lrwEph+EifxSJqG3jACZLASSBUcuHGq/2ZgyM=",
+      "flash_bwd_preprocess.py": "A5/CJhk0Zg6mMb0PU4Lk9mwjDhr3/WaxlFaBtNOhBac=",
+      "flash_dec_combine.py": "7ZHMacWY+QPPjgSezKmQzhogibUZfb7abIwQW+2U6ak=",
+      "flash_dense_bwd.py": "sSRrZs5y0VB5Ga/9KQCAmuYaKwxeoVJfhwrumzby7xw=",
+      "flash_dense_dec.py": "9/JfVxX+4hz3Bq/4liJih6pDwjEZiOyHAZeUYYlMyjM=",
+      "flash_dense_fwd.py": "ydJnTHmzMX52WOyGilqNf7pj+6qimeCjc6NmSDJs9BY=",
+      "flash_fwd_combine.py": "8iYebjXMgL3fWzR1hiA6/PNLVLjCnze6PmgrI+D/cOA=",
+      "flash_gated_bwd.py": "i9+fAUdp0xxpk7A/V6KnLwMVThdrOlFpaWwurQDXyDU=",
+      "flash_gated_dec.py": "fLJHd8ii4lV5SCcLYvuDwZVT3Q+x7ANgsQ3xojZgHrw=",
+      "flash_gated_fwd.py": "W/9zB7NdC8Ap/dNuMYSNGMMaTxpyQTeyUk2AAKL8vvU=",
+      "flash_proj_bwd.py": "a+CpQe6DsiNeVZlq+AJDTjdSeqquUAhege37vrvckDo=",
+      "flash_proj_fwd.py": "TVuL7nkXcTi+R5j1tTd5gwFPWl1ddbvnupER5QhUQtY=",
+      "flash_sparse_attention/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
+      "flash_sparse_bwd.py": "sPh3T1XiomcbSHb/+fUvYZBYgWZNCuwMJW5UoabQUPE=",
+      "flash_sparse_dec.py": "yqOVdw3aVkSoD1aDtqng1bmNn3eWQkTLphVRm+3vBOA=",
+      "flash_sparse_fwd.py": "7cUeJrayY/yu/0+FpTqG2foJVfdweaq7KuZsge7/rno=",
+      "interface.py": "0b2Ya0TXU/6QS5plzsK6FOStUJNGl49fH7EsspTMkoo=",
+      "kernel_repr.py": "GdLRDHO7bMvfWt52vCDqO9HLxQeaK7/sr3iMeeSTJnA=",
+      "launch_grid.py": "NKGkqIro3SDRikGl1C7Om0/2DVH0eN8ecepkd6h8A98=",
+      "launch_template.py": "CbrryDz3BPY8zpWcTo6hNpI0DKfW16BkY66f0oNNAYs=",
+      "mask.py": "aF/KiRajGsdsLdQllFO99NBrOQeTjCPK9aCIunxen/0=",
+      "quant.py": "B47A0bWsIb5E5dynB0NrcJjg1l1F/rj/7y6yVXIWrNI=",
+      "scheduler.py": "4sgkUQyAuijIyfIk88/kyD4a+dI2LWeptB28Pfe2nko=",
+      "seqlen_info.py": "BCd/B9qu9DfS2+VxX9uBdh8qxrdDYgjJzwfp7Q7i6Bc=",
+      "utils.py": "g8U91hmiF+yFwiGEVpYzq1k1CMH+3sIrHq0nrOaNrMI="
+    }
+  }
+}

build/torch-cuda/quant.py ADDED Viewed

	@@ -0,0 +1,66 @@

+import torch
+import triton
+import triton.language as tl
+@triton.jit
+def _amax_kernel(
+    X,
+    Amax,
+    n_elements,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offsets = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offsets < n_elements
+    x = tl.load(X + offsets, mask=mask, other=0.0, cache_modifier=".cg")
+    local_amax = tl.max(tl.abs(x))
+    tl.atomic_max(Amax, local_amax, sem="relaxed")
+@triton.jit
+def _quantize_fp8_kernel(
+    X,
+    Out,
+    Scale,
+    n_elements,
+    BLOCK: tl.constexpr,
+    FP8_DTYPE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offsets = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offsets < n_elements
+    x = tl.load(X + offsets, mask=mask, other=0.0, cache_modifier=".cg")
+    scale = tl.load(Scale)
+    x_fp8 = (x * scale).to(FP8_DTYPE)
+    tl.store(Out + offsets, x_fp8, mask=mask)
+def quantize_fp8(x: torch.Tensor, dtype: torch.dtype = torch.float8_e5m2):
+    assert x.is_cuda and x.dtype in (torch.float16, torch.bfloat16, torch.float32)
+    assert dtype in (torch.float8_e5m2, torch.float8_e4m3fn)
+    tl_dtype = tl.float8e5 if dtype == torch.float8_e5m2 else tl.float8e4nv
+    FP8_MAX = 57344.0 if dtype == torch.float8_e5m2 else 448.0
+    n_elements = x.numel()
+    amax = torch.zeros(1, device=x.device, dtype=torch.float32)
+    x_fp8 = torch.empty_like(x, dtype=dtype)
+    BLOCK = 512
+    def grid(META):
+        return (triton.cdiv(n_elements, META["BLOCK"]),)
+    _amax_kernel[grid](x, amax, n_elements, BLOCK=BLOCK)
+    scale = FP8_MAX / amax
+    _quantize_fp8_kernel[grid](
+        x, x_fp8, scale, n_elements, BLOCK=BLOCK, FP8_DTYPE=tl_dtype
+    )
+    return x_fp8, (amax / FP8_MAX).to(x.dtype)

build/torch-cuda/scheduler.py ADDED Viewed

	@@ -0,0 +1,2609 @@

+import triton
+import triton.language as tl
+from triton.language.core import _aggregate as aggregate
+from triton.runtime.jit import constexpr_function
+from . import mask
+from . import activations
+from . import seqlen_info, block_info
+@aggregate
+class AttnFwdGridIndex:
+    m_block: tl.tensor
+    batch_idx: tl.tensor
+    head_idx: tl.tensor
+    head_kv_idx: tl.tensor
+    split_idx: tl.tensor
+    @constexpr_function
+    def __init__(self, m_block, batch_idx, head_idx, head_kv_idx, split_idx):
+        self.m_block = m_block
+        self.batch_idx = batch_idx
+        self.head_idx = head_idx
+        self.head_kv_idx = head_kv_idx
+        self.split_idx = split_idx
+    @staticmethod
+    @triton.jit
+    def create(
+        num_splits,
+        QHEAD_PER_KVHEAD: tl.constexpr,
+        IS_SPLIT_KV: tl.constexpr,
+        PACK_GQA: tl.constexpr,
+    ):
+        m_block = tl.program_id(0)
+        head_idx = tl.program_id(1)
+        batch_split_idx = tl.program_id(2)
+        if IS_SPLIT_KV:
+            batch_idx = batch_split_idx // num_splits
+            split_idx = batch_split_idx - batch_idx * num_splits
+        else:
+            batch_idx = batch_split_idx
+            split_idx = 0
+        if PACK_GQA:
+            head_kv_idx = head_idx
+        else:
+            head_kv_idx = head_idx // QHEAD_PER_KVHEAD
+        return AttnFwdGridIndex(m_block, batch_idx, head_idx, head_kv_idx, split_idx)
+    @triton.jit
+    def load_window_sizes(self, window_sizes, stride_wh, IS_LOCAL: tl.constexpr):
+        if IS_LOCAL:
+            window_size_sink = tl.load(window_sizes + self.head_kv_idx * stride_wh)
+            window_size_left = tl.load(window_sizes + self.head_kv_idx * stride_wh + 1)
+            window_size_right = tl.load(window_sizes + self.head_kv_idx * stride_wh + 2)
+            window_size_dist = tl.load(window_sizes + self.head_kv_idx * stride_wh + 3)
+        else:
+            window_size_sink = 0
+            window_size_left = 0
+            window_size_right = 0
+            window_size_dist = 0
+        return window_size_sink, window_size_left, window_size_right, window_size_dist
+@aggregate
+class AttnBwdGridIndex:
+    n_block: tl.tensor
+    batch_idx: tl.tensor
+    head_idx: tl.tensor
+    head_kv_idx: tl.tensor
+    split_idx: tl.tensor
+    @constexpr_function
+    def __init__(self, n_block, batch_idx, head_idx, head_kv_idx, split_idx):
+        self.n_block = n_block
+        self.batch_idx = batch_idx
+        self.head_idx = head_idx
+        self.head_kv_idx = head_kv_idx
+        self.split_idx = split_idx
+    @staticmethod
+    @triton.jit
+    def create(
+        num_splits,
+        QHEAD_PER_KVHEAD: tl.constexpr,
+        IS_SPLIT_QO: tl.constexpr,
+    ):
+        n_block = tl.program_id(0)
+        head_idx = tl.program_id(1)
+        batch_split_idx = tl.program_id(2)
+        if IS_SPLIT_QO:
+            batch_idx = batch_split_idx // num_splits
+            split_idx = batch_split_idx - batch_idx * num_splits
+        else:
+            batch_idx = batch_split_idx
+            split_idx = 0
+        head_kv_idx = head_idx // QHEAD_PER_KVHEAD
+        return AttnBwdGridIndex(n_block, batch_idx, head_idx, head_kv_idx, split_idx)
+    @triton.jit
+    def load_window_sizes(self, window_sizes, stride_wh, IS_LOCAL: tl.constexpr):
+        if IS_LOCAL:
+            window_size_sink = tl.load(window_sizes + self.head_kv_idx * stride_wh)
+            window_size_left = tl.load(window_sizes + self.head_kv_idx * stride_wh + 1)
+            window_size_right = tl.load(window_sizes + self.head_kv_idx * stride_wh + 2)
+            window_size_dist = tl.load(window_sizes + self.head_kv_idx * stride_wh + 3)
+        else:
+            window_size_sink = 0
+            window_size_left = 0
+            window_size_right = 0
+            window_size_dist = 0
+        return window_size_sink, window_size_left, window_size_right, window_size_dist
+@aggregate
+class AttnDecGridIndex:
+    batch_idx: tl.tensor
+    head_idx: tl.tensor
+    head_kv_idx: tl.tensor
+    split_idx: tl.tensor
+    @constexpr_function
+    def __init__(self, batch_idx, head_idx, head_kv_idx, split_idx):
+        self.batch_idx = batch_idx
+        self.head_idx = head_idx
+        self.head_kv_idx = head_kv_idx
+        self.split_idx = split_idx
+    @staticmethod
+    @triton.jit
+    def create(
+        num_splits,
+    ):
+        head_idx = tl.program_id(0)
+        batch_split_idx = tl.program_id(1)
+        batch_idx = batch_split_idx // num_splits
+        split_idx = batch_split_idx - batch_idx * num_splits
+        head_kv_idx = head_idx
+        return AttnDecGridIndex(batch_idx, head_idx, head_kv_idx, split_idx)
+    @triton.jit
+    def load_window_sizes(self, window_sizes, stride_wh, IS_LOCAL: tl.constexpr):
+        if IS_LOCAL:
+            window_size_sink = tl.load(window_sizes + self.head_kv_idx * stride_wh)
+            window_size_left = tl.load(window_sizes + self.head_kv_idx * stride_wh + 1)
+            window_size_right = tl.load(window_sizes + self.head_kv_idx * stride_wh + 2)
+            window_size_dist = tl.load(window_sizes + self.head_kv_idx * stride_wh + 3)
+        else:
+            window_size_sink = 0
+            window_size_left = 0
+            window_size_right = 0
+            window_size_dist = 0
+        return window_size_sink, window_size_left, window_size_right, window_size_dist
+@aggregate
+class AttnFwdConfig:
+    softmax_scale_log2: tl.tensor
+    softmax_threshold_log2: tl.tensor
+    gate_threshold_log2: tl.tensor
+    value_scale: tl.tensor
+    m_block: tl.tensor
+    actual_seqlen_q: tl.tensor
+    actual_seqlen_k: tl.tensor
+    offset_q: tl.tensor
+    offset_k: tl.tensor
+    padded_offset_q: tl.tensor
+    padded_offset_k: tl.tensor
+    window_size_sink: tl.tensor
+    window_size_left: tl.tensor
+    window_size_right: tl.tensor
+    window_size_dist: tl.tensor
+    head_dim: tl.tensor
+    PACK_GQA: tl.constexpr
+    QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr
+    TILE_M: tl.constexpr
+    TILE_N: tl.constexpr
+    TILE_K: tl.constexpr
+    IS_LOGSIGMOID_GATE: tl.constexpr
+    @constexpr_function
+    def __init__(
+        self,
+        softmax_scale_log2,
+        softmax_threshold_log2,
+        gate_threshold_log2,
+        value_scale,
+        m_block,
+        actual_seqlen_q,
+        actual_seqlen_k,
+        offset_q,
+        offset_k,
+        padded_offset_q,
+        padded_offset_k,
+        window_size_sink,
+        window_size_left,
+        window_size_right,
+        window_size_dist,
+        head_dim,
+        PACK_GQA,
+        QHEAD_PER_KVHEAD_PACKGQA,
+        TILE_M,
+        TILE_N,
+        TILE_K,
+        IS_LOGSIGMOID_GATE,
+    ):
+        self.softmax_scale_log2 = softmax_scale_log2
+        self.softmax_threshold_log2 = softmax_threshold_log2
+        self.gate_threshold_log2 = gate_threshold_log2
+        self.value_scale = value_scale
+        self.m_block = m_block
+        self.actual_seqlen_q = actual_seqlen_q
+        self.actual_seqlen_k = actual_seqlen_k
+        self.offset_q = offset_q
+        self.offset_k = offset_k
+        self.padded_offset_q = padded_offset_q
+        self.padded_offset_k = padded_offset_k
+        self.window_size_sink = window_size_sink
+        self.window_size_left = window_size_left
+        self.window_size_right = window_size_right
+        self.window_size_dist = window_size_dist
+        self.head_dim = head_dim
+        self.PACK_GQA = tl.constexpr(PACK_GQA)
+        self.QHEAD_PER_KVHEAD_PACKGQA = tl.constexpr(QHEAD_PER_KVHEAD_PACKGQA)
+        self.TILE_M = tl.constexpr(TILE_M)
+        self.TILE_N = tl.constexpr(TILE_N)
+        self.TILE_K = tl.constexpr(TILE_K)
+        self.IS_LOGSIGMOID_GATE = tl.constexpr(IS_LOGSIGMOID_GATE)
+    @staticmethod
+    @triton.jit
+    def create(
+        softmax_scale=0.0,
+        softmax_threshold=0.0,
+        gate_threshold=0.0,
+        query_scale=None,
+        key_scale=None,
+        value_scale=None,
+        m_block=0,
+        batch_idx=0,
+        window_size_sink=0,
+        window_size_left=0,
+        window_size_right=0,
+        window_size_dist=0,
+        head_dim=0,
+        cu_seqlens_q=None,
+        cu_seqlens_k=None,
+        seqused_q=None,
+        seqused_k=None,
+        seqlen_q=0,
+        seqlen_k=0,
+        PACK_GQA: tl.constexpr = False,
+        QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr = 1,
+        TILE_M: tl.constexpr = 64,
+        TILE_N: tl.constexpr = 64,
+        TILE_K: tl.constexpr = 64,
+        IS_CAUSAL: tl.constexpr = False,
+        IS_LOGSIGMOID_GATE: tl.constexpr = False,
+        IS_ADAPT_GATE: tl.constexpr = False,
+        HAS_CU_SEQLENS_Q: tl.constexpr = False,
+        HAS_CU_SEQLENS_K: tl.constexpr = False,
+        HAS_SEQUSED_Q: tl.constexpr = False,
+        HAS_SEQUSED_K: tl.constexpr = False,
+    ):
+        # Get seqlen info for this batch
+        (
+            offset_q,
+            offset_k,
+            padded_offset_q,
+            padded_offset_k,
+            actual_seqlen_q,
+            actual_seqlen_k,
+        ) = seqlen_info.get_seqlen_info_qk(
+            batch_idx=batch_idx,
+            seqlen_q_static=seqlen_q,
+            seqlen_k_static=seqlen_k,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+            TILE_M=TILE_M,
+            TILE_N=TILE_N,
+            HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+            HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+            HAS_SEQUSED_Q=HAS_SEQUSED_Q,
+            HAS_SEQUSED_K=HAS_SEQUSED_K,
+        )
+        # Load query scale
+        q_scale = tl.load(query_scale)
+        # Load key scale
+        k_scale = tl.load(key_scale)
+        # Load value scale
+        v_scale = tl.load(value_scale)
+        LOG2E: tl.constexpr = 1.44269504089
+        softmax_scale_log2 = softmax_scale * LOG2E * q_scale * k_scale
+        softmax_threshold_log2 = seqlen_info.get_softmax_threshold(
+            softmax_threshold,
+            m_block,
+            actual_seqlen_q,
+            actual_seqlen_k,
+            IS_CAUSAL,
+            TILE_M,
+            QHEAD_PER_KVHEAD_PACKGQA,
+        )
+        gate_threshold_log2 = seqlen_info.get_gate_threshold(
+            gate_threshold,
+            m_block,
+            actual_seqlen_q,
+            actual_seqlen_k,
+            IS_CAUSAL,
+            TILE_M,
+            QHEAD_PER_KVHEAD_PACKGQA,
+            IS_ADAPT_GATE,
+        )
+        return AttnFwdConfig(
+            softmax_scale_log2,
+            softmax_threshold_log2,
+            gate_threshold_log2,
+            v_scale,
+            m_block,
+            actual_seqlen_q,
+            actual_seqlen_k,
+            offset_q,
+            offset_k,
+            padded_offset_q,
+            padded_offset_k,
+            window_size_sink,
+            window_size_left,
+            window_size_right,
+            window_size_dist,
+            head_dim,
+            PACK_GQA,
+            QHEAD_PER_KVHEAD_PACKGQA,
+            TILE_M,
+            TILE_N,
+            TILE_K,
+            IS_LOGSIGMOID_GATE,
+        )
+    @triton.jit
+    def get_offs_m(self):
+        return self.m_block * self.TILE_M + tl.arange(0, self.TILE_M)
+    @triton.jit
+    def get_offs_k(self):
+        return tl.arange(0, self.TILE_K)
+@aggregate
+class AttnBwdConfig:
+    softmax_scale_log2: tl.tensor
+    softmax_threshold: tl.tensor
+    gate_threshold: tl.tensor
+    query_scale: tl.tensor
+    key_scale: tl.tensor
+    value_scale: tl.tensor
+    n_block: tl.tensor
+    actual_seqlen_q: tl.tensor
+    actual_seqlen_k: tl.tensor
+    offset_q: tl.tensor
+    offset_k: tl.tensor
+    padded_offset_q: tl.tensor
+    padded_offset_k: tl.tensor
+    window_size_sink: tl.tensor
+    window_size_left: tl.tensor
+    window_size_right: tl.tensor
+    window_size_dist: tl.tensor
+    head_dim: tl.tensor
+    QHEAD_PER_KVHEAD: tl.constexpr
+    TILE_M: tl.constexpr
+    TILE_N: tl.constexpr
+    TILE_K: tl.constexpr
+    IS_CAUSAL: tl.constexpr
+    IS_LOGSIGMOID_GATE: tl.constexpr
+    IS_ADAPT_GATE: tl.constexpr
+    @constexpr_function
+    def __init__(
+        self,
+        softmax_scale_log2,
+        softmax_threshold,
+        gate_threshold,
+        query_scale,
+        key_scale,
+        value_scale,
+        n_block,
+        actual_seqlen_q,
+        actual_seqlen_k,
+        offset_q,
+        offset_k,
+        padded_offset_q,
+        padded_offset_k,
+        window_size_sink,
+        window_size_left,
+        window_size_right,
+        window_size_dist,
+        head_dim,
+        QHEAD_PER_KVHEAD,
+        TILE_M,
+        TILE_N,
+        TILE_K,
+        IS_CAUSAL,
+        IS_LOGSIGMOID_GATE,
+        IS_ADAPT_GATE,
+    ):
+        self.softmax_scale_log2 = softmax_scale_log2
+        self.softmax_threshold = softmax_threshold
+        self.gate_threshold = gate_threshold
+        self.query_scale = query_scale
+        self.key_scale = key_scale
+        self.value_scale = value_scale
+        self.n_block = n_block
+        self.actual_seqlen_q = actual_seqlen_q
+        self.actual_seqlen_k = actual_seqlen_k
+        self.offset_q = offset_q
+        self.offset_k = offset_k
+        self.padded_offset_q = padded_offset_q
+        self.padded_offset_k = padded_offset_k
+        self.window_size_sink = window_size_sink
+        self.window_size_left = window_size_left
+        self.window_size_right = window_size_right
+        self.window_size_dist = window_size_dist
+        self.head_dim = head_dim
+        self.QHEAD_PER_KVHEAD = tl.constexpr(QHEAD_PER_KVHEAD)
+        self.TILE_M = tl.constexpr(TILE_M)
+        self.TILE_N = tl.constexpr(TILE_N)
+        self.TILE_K = tl.constexpr(TILE_K)
+        self.IS_CAUSAL = tl.constexpr(IS_CAUSAL)
+        self.IS_LOGSIGMOID_GATE = tl.constexpr(IS_LOGSIGMOID_GATE)
+        self.IS_ADAPT_GATE = tl.constexpr(IS_ADAPT_GATE)
+    @staticmethod
+    @triton.jit
+    def create(
+        softmax_scale=0.0,
+        softmax_threshold=0.0,
+        gate_threshold=0.0,
+        query_scale=None,
+        key_scale=None,
+        value_scale=None,
+        n_block=0,
+        batch_idx=0,
+        window_size_sink=0,
+        window_size_left=0,
+        window_size_right=0,
+        window_size_dist=0,
+        head_dim=0,
+        cu_seqlens_q=None,
+        cu_seqlens_k=None,
+        seqused_q=None,
+        seqused_k=None,
+        seqlen_q=0,
+        seqlen_k=0,
+        QHEAD_PER_KVHEAD: tl.constexpr = 1,
+        TILE_M: tl.constexpr = 64,
+        TILE_N: tl.constexpr = 64,
+        TILE_K: tl.constexpr = 64,
+        IS_CAUSAL: tl.constexpr = False,
+        IS_LOGSIGMOID_GATE: tl.constexpr = False,
+        IS_ADAPT_GATE: tl.constexpr = False,
+        HAS_CU_SEQLENS_Q: tl.constexpr = False,
+        HAS_CU_SEQLENS_K: tl.constexpr = False,
+        HAS_SEQUSED_Q: tl.constexpr = False,
+        HAS_SEQUSED_K: tl.constexpr = False,
+    ):
+        # Get seqlen info for this batch
+        (
+            offset_q,
+            offset_k,
+            padded_offset_q,
+            padded_offset_k,
+            actual_seqlen_q,
+            actual_seqlen_k,
+        ) = seqlen_info.get_seqlen_info_qk(
+            batch_idx=batch_idx,
+            seqlen_q_static=seqlen_q,
+            seqlen_k_static=seqlen_k,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+            TILE_M=TILE_M,
+            TILE_N=TILE_N,
+            HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+            HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+            HAS_SEQUSED_Q=HAS_SEQUSED_Q,
+            HAS_SEQUSED_K=HAS_SEQUSED_K,
+        )
+        # Load query scale
+        q_scale = tl.load(query_scale)
+        # Load key scale
+        k_scale = tl.load(key_scale)
+        # Load value scale
+        v_scale = tl.load(value_scale)
+        LOG2E: tl.constexpr = 1.44269504089
+        softmax_scale_log2 = softmax_scale * LOG2E
+        softmax_threshold = tl.full((), softmax_threshold, tl.float32)
+        gate_threshold = tl.full((), gate_threshold, tl.float32)
+        return AttnBwdConfig(
+            softmax_scale_log2,
+            softmax_threshold,
+            gate_threshold,
+            q_scale,
+            k_scale,
+            v_scale,
+            n_block,
+            actual_seqlen_q,
+            actual_seqlen_k,
+            offset_q,
+            offset_k,
+            padded_offset_q,
+            padded_offset_k,
+            window_size_sink,
+            window_size_left,
+            window_size_right,
+            window_size_dist,
+            head_dim,
+            QHEAD_PER_KVHEAD,
+            TILE_M,
+            TILE_N,
+            TILE_K,
+            IS_CAUSAL,
+            IS_LOGSIGMOID_GATE,
+            IS_ADAPT_GATE,
+        )
+    @triton.jit
+    def get_softmax_threshold_log2(
+        self,
+        m_block,
+    ):
+        return seqlen_info.get_softmax_threshold(
+            self.softmax_threshold,
+            m_block,
+            self.actual_seqlen_q,
+            self.actual_seqlen_k,
+            self.IS_CAUSAL,
+            self.TILE_M,
+            QHEAD_PER_KVHEAD_PACKGQA=1,
+        )
+    @triton.jit
+    def get_gate_threshold_log2(
+        self,
+        m_block,
+    ):
+        return seqlen_info.get_gate_threshold(
+            self.gate_threshold,
+            m_block,
+            self.actual_seqlen_q,
+            self.actual_seqlen_k,
+            self.IS_CAUSAL,
+            self.TILE_M,
+            QHEAD_PER_KVHEAD_PACKGQA=1,
+            IS_ADAPT_GATE=self.IS_ADAPT_GATE,
+        )
+@aggregate
+class AttnDecConfig:
+    softmax_scale_log2: tl.tensor
+    softmax_threshold_log2: tl.tensor
+    gate_threshold_log2: tl.tensor
+    value_scale: tl.tensor
+    m_block: tl.tensor
+    actual_seqlen_q: tl.tensor
+    actual_seqlen_k: tl.tensor
+    offset_q: tl.tensor
+    offset_k: tl.tensor
+    padded_offset_q: tl.tensor
+    padded_offset_k: tl.tensor
+    window_size_sink: tl.tensor
+    window_size_left: tl.tensor
+    window_size_right: tl.tensor
+    window_size_dist: tl.tensor
+    head_dim: tl.tensor
+    QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr
+    TILE_M: tl.constexpr
+    TILE_N: tl.constexpr
+    TILE_K: tl.constexpr
+    IS_LOGSIGMOID_GATE: tl.constexpr
+    @constexpr_function
+    def __init__(
+        self,
+        softmax_scale_log2,
+        softmax_threshold_log2,
+        gate_threshold_log2,
+        value_scale,
+        m_block,
+        actual_seqlen_q,
+        actual_seqlen_k,
+        offset_q,
+        offset_k,
+        padded_offset_q,
+        padded_offset_k,
+        window_size_sink,
+        window_size_left,
+        window_size_right,
+        window_size_dist,
+        head_dim,
+        QHEAD_PER_KVHEAD_PACKGQA,
+        TILE_M,
+        TILE_N,
+        TILE_K,
+        IS_LOGSIGMOID_GATE,
+    ):
+        self.softmax_scale_log2 = softmax_scale_log2
+        self.softmax_threshold_log2 = softmax_threshold_log2
+        self.gate_threshold_log2 = gate_threshold_log2
+        self.value_scale = value_scale
+        self.m_block = m_block
+        self.actual_seqlen_q = actual_seqlen_q
+        self.actual_seqlen_k = actual_seqlen_k
+        self.offset_q = offset_q
+        self.offset_k = offset_k
+        self.padded_offset_q = padded_offset_q
+        self.padded_offset_k = padded_offset_k
+        self.window_size_sink = window_size_sink
+        self.window_size_left = window_size_left
+        self.window_size_right = window_size_right
+        self.window_size_dist = window_size_dist
+        self.head_dim = head_dim
+        self.QHEAD_PER_KVHEAD_PACKGQA = tl.constexpr(QHEAD_PER_KVHEAD_PACKGQA)
+        self.TILE_M = tl.constexpr(TILE_M)
+        self.TILE_N = tl.constexpr(TILE_N)
+        self.TILE_K = tl.constexpr(TILE_K)
+        self.IS_LOGSIGMOID_GATE = tl.constexpr(IS_LOGSIGMOID_GATE)
+    @staticmethod
+    @triton.jit
+    def create(
+        softmax_scale=0.0,
+        softmax_threshold=0.0,
+        gate_threshold=0.0,
+        query_scale=None,
+        key_scale=None,
+        value_scale=None,
+        batch_idx=0,
+        window_size_sink=0,
+        window_size_left=0,
+        window_size_right=0,
+        window_size_dist=0,
+        head_dim=0,
+        cu_seqlens_q=None,
+        cu_seqlens_k=None,
+        seqused_q=None,
+        seqused_k=None,
+        seqlen_q=0,
+        seqlen_k=0,
+        QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr = 1,
+        TILE_M: tl.constexpr = 16,
+        TILE_N: tl.constexpr = 64,
+        TILE_K: tl.constexpr = 64,
+        IS_LOGSIGMOID_GATE: tl.constexpr = False,
+        HAS_CU_SEQLENS_Q: tl.constexpr = False,
+        HAS_CU_SEQLENS_K: tl.constexpr = False,
+        HAS_SEQUSED_Q: tl.constexpr = False,
+        HAS_SEQUSED_K: tl.constexpr = False,
+    ):
+        # Get seqlen info for this batch
+        (
+            offset_q,
+            offset_k,
+            padded_offset_q,
+            padded_offset_k,
+            actual_seqlen_q,
+            actual_seqlen_k,
+        ) = seqlen_info.get_seqlen_info_qk(
+            batch_idx=batch_idx,
+            seqlen_q_static=seqlen_q,
+            seqlen_k_static=seqlen_k,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+            TILE_M=TILE_M,
+            TILE_N=TILE_N,
+            HAS_CU_SEQLENS_Q=HAS_CU_SEQLENS_Q,
+            HAS_CU_SEQLENS_K=HAS_CU_SEQLENS_K,
+            HAS_SEQUSED_Q=HAS_SEQUSED_Q,
+            HAS_SEQUSED_K=HAS_SEQUSED_K,
+        )
+        # Load query scale
+        q_scale = tl.load(query_scale)
+        # Load key scale
+        k_scale = tl.load(key_scale)
+        # Load value scale
+        v_scale = tl.load(value_scale)
+        LOG2E: tl.constexpr = 1.44269504089
+        softmax_scale_log2 = softmax_scale * LOG2E * q_scale * k_scale
+        softmax_threshold_log2 = seqlen_info.get_softmax_threshold(
+            softmax_threshold,
+            0,
+            actual_seqlen_q,
+            actual_seqlen_k,
+            False,
+            TILE_M,
+            QHEAD_PER_KVHEAD_PACKGQA,
+        )
+        gate_threshold_log2 = seqlen_info.get_gate_threshold(
+            gate_threshold,
+            0,
+            actual_seqlen_q,
+            actual_seqlen_k,
+            False,
+            TILE_M,
+            QHEAD_PER_KVHEAD_PACKGQA,
+            False,
+        )
+        return AttnDecConfig(
+            softmax_scale_log2,
+            softmax_threshold_log2,
+            gate_threshold_log2,
+            v_scale,
+            tl.full((), 0, tl.int32),
+            actual_seqlen_q,
+            actual_seqlen_k,
+            offset_q,
+            offset_k,
+            padded_offset_q,
+            padded_offset_k,
+            window_size_sink,
+            window_size_left,
+            window_size_right,
+            window_size_dist,
+            head_dim,
+            QHEAD_PER_KVHEAD_PACKGQA,
+            TILE_M,
+            TILE_N,
+            TILE_K,
+            IS_LOGSIGMOID_GATE,
+        )
+@aggregate
+class AttnFwdBlockScheduler:
+    n_block_min: tl.tensor
+    n_block_max: tl.tensor
+    n_block_max_no_mask: tl.tensor
+    n_block_window_min: tl.tensor
+    n_block_window_max: tl.tensor
+    n_block_window_min_no_mask: tl.tensor
+    n_block_window_max_no_mask: tl.tensor
+    n_block_sink_min: tl.tensor
+    n_block_sink_max: tl.tensor
+    @constexpr_function
+    def __init__(
+        self,
+        n_block_min,
+        n_block_max,
+        n_block_max_no_mask,
+        n_block_window_min,
+        n_block_window_max,
+        n_block_window_min_no_mask,
+        n_block_window_max_no_mask,
+        n_block_sink_min,
+        n_block_sink_max,
+    ):
+        self.n_block_min = n_block_min
+        self.n_block_max = n_block_max
+        self.n_block_max_no_mask = n_block_max_no_mask
+        self.n_block_window_min = n_block_window_min
+        self.n_block_window_max = n_block_window_max
+        self.n_block_window_min_no_mask = n_block_window_min_no_mask
+        self.n_block_window_max_no_mask = n_block_window_max_no_mask
+        self.n_block_sink_min = n_block_sink_min
+        self.n_block_sink_max = n_block_sink_max
+    @triton.jit
+    def is_empty(self):
+        return (
+            (self.n_block_max <= self.n_block_min)
+            & (self.n_block_window_max <= self.n_block_window_min)
+            & (self.n_block_sink_max <= self.n_block_sink_min)
+        )
+    @staticmethod
+    @triton.jit
+    def create(
+        config: AttnFwdConfig,
+        split_idx,
+        num_splits,
+        IS_CAUSAL: tl.constexpr,
+        IS_LOCAL: tl.constexpr,
+        IS_SPLIT_KV: tl.constexpr,
+    ):
+        # Compute causal n_block range for this m_block
+        (
+            n_block_min,
+            n_block_max,
+            n_block_window_min,
+            n_block_window_max,
+            n_block_sink_min,
+            n_block_sink_max,
+        ) = block_info.get_n_block_min_max(
+            seqlen_q=config.actual_seqlen_q,
+            seqlen_k=config.actual_seqlen_k,
+            m_block=config.m_block,
+            split_idx=split_idx,
+            num_splits=num_splits,
+            window_size_sink=config.window_size_sink,
+            window_size_left=config.window_size_left,
+            window_size_right=config.window_size_right,
+            window_size_dist=config.window_size_dist,
+            TILE_N=config.TILE_N,
+            TILE_M=config.TILE_M,
+            IS_CAUSAL=IS_CAUSAL,
+            IS_LOCAL=IS_LOCAL,
+            IS_SPLIT_KV=IS_SPLIT_KV,
+            QHEAD_PER_KVHEAD_PACKGQA=config.QHEAD_PER_KVHEAD_PACKGQA,
+        )
+        n_block_max_no_mask = block_info.get_n_block_min_causal_local_mask(
+            seqlen_q=config.actual_seqlen_q,
+            seqlen_k=config.actual_seqlen_k,
+            m_block=config.m_block,
+            n_block_min=n_block_min,
+            window_size_right=0,
+            window_size_dist=0,
+            TILE_N=config.TILE_N,
+            TILE_M=config.TILE_M,
+            IS_LOCAL=False,
+            QHEAD_PER_KVHEAD_PACKGQA=config.QHEAD_PER_KVHEAD_PACKGQA,
+        )
+        # Clamp to split's range so the no-mask loop stays within bounds
+        if IS_SPLIT_KV:
+            n_block_max_no_mask = tl.minimum(n_block_max_no_mask, n_block_max)
+        if IS_LOCAL:
+            # Compute local n_block range for this m_block
+            if IS_SPLIT_KV:
+                n_block_max_no_mask = tl.where(
+                    split_idx >= num_splits - 1,
+                    n_block_min,
+                    n_block_max_no_mask,
+                )
+            else:
+                n_block_max_no_mask = n_block_min
+            if IS_SPLIT_KV:
+                n_block_window_min = tl.maximum(n_block_window_min, n_block_min)
+            n_block_window_max = tl.minimum(n_block_window_max, n_block_max_no_mask)
+            n_block_window_max_no_mask = block_info.get_n_block_min_causal_local_mask(
+                seqlen_q=config.actual_seqlen_q,
+                seqlen_k=config.actual_seqlen_k,
+                m_block=config.m_block,
+                n_block_min=n_block_window_min,
+                window_size_right=config.window_size_right,
+                window_size_dist=config.window_size_dist,
+                TILE_N=config.TILE_N,
+                TILE_M=config.TILE_M,
+                IS_LOCAL=True,
+                QHEAD_PER_KVHEAD_PACKGQA=config.QHEAD_PER_KVHEAD_PACKGQA,
+            )
+            n_block_window_min_no_mask = block_info.get_n_block_min_before_local_mask(
+                seqlen_q=config.actual_seqlen_q,
+                seqlen_k=config.actual_seqlen_k,
+                m_block=config.m_block,
+                n_block_min=n_block_window_min,
+                window_size_left=config.window_size_left,
+                window_size_right=config.window_size_right,
+                window_size_dist=config.window_size_dist,
+                TILE_N=config.TILE_N,
+                TILE_M=config.TILE_M,
+                IS_LOCAL=True,
+                QHEAD_PER_KVHEAD_PACKGQA=config.QHEAD_PER_KVHEAD_PACKGQA,
+            )
+            n_block_window_min_no_mask = tl.minimum(
+                n_block_window_min_no_mask, n_block_window_max_no_mask
+            )
+            # Clamp window no-mask boundaries to the window's range
+            if IS_SPLIT_KV:
+                n_block_window_max_no_mask = tl.maximum(
+                    tl.minimum(n_block_window_max_no_mask, n_block_window_max),
+                    n_block_window_min,
+                )
+                n_block_window_min_no_mask = tl.maximum(
+                    tl.minimum(n_block_window_min_no_mask, n_block_window_max),
+                    n_block_window_min,
+                )
+        else:
+            n_block_window_min = 0
+            n_block_window_max = 0
+            n_block_window_min_no_mask = 0
+            n_block_window_max_no_mask = 0
+        return AttnFwdBlockScheduler(
+            n_block_min,
+            n_block_max,
+            n_block_max_no_mask,
+            n_block_window_min,
+            n_block_window_max,
+            n_block_window_min_no_mask,
+            n_block_window_max_no_mask,
+            n_block_sink_min,
+            n_block_sink_max,
+        )
+@aggregate
+class AttnBwdBlockScheduler:
+    m_block_min: tl.tensor
+    m_block_max: tl.tensor
+    m_block_min_no_mask: tl.tensor
+    m_block_window_min: tl.tensor
+    m_block_window_max: tl.tensor
+    m_block_window_min_no_mask: tl.tensor
+    m_block_window_max_no_mask: tl.tensor
+    m_block_sink_min: tl.tensor
+    m_block_sink_max: tl.tensor
+    @constexpr_function
+    def __init__(
+        self,
+        m_block_min,
+        m_block_max,
+        m_block_min_no_mask,
+        m_block_window_min,
+        m_block_window_max,
+        m_block_window_min_no_mask,
+        m_block_window_max_no_mask,
+        m_block_sink_min,
+        m_block_sink_max,
+    ):
+        self.m_block_min = m_block_min
+        self.m_block_max = m_block_max
+        self.m_block_min_no_mask = m_block_min_no_mask
+        self.m_block_window_min = m_block_window_min
+        self.m_block_window_max = m_block_window_max
+        self.m_block_window_min_no_mask = m_block_window_min_no_mask
+        self.m_block_window_max_no_mask = m_block_window_max_no_mask
+        self.m_block_sink_min = m_block_sink_min
+        self.m_block_sink_max = m_block_sink_max
+    @triton.jit
+    def is_empty(self):
+        return (
+            (self.m_block_max <= self.m_block_min)
+            & (self.m_block_window_max <= self.m_block_window_min)
+            & (self.m_block_sink_max <= self.m_block_sink_min)
+        )
+    @staticmethod
+    @triton.jit
+    def create(
+        config: AttnBwdConfig,
+        split_idx,
+        num_splits,
+        IS_CAUSAL: tl.constexpr,
+        IS_LOCAL: tl.constexpr,
+        IS_SPLIT_QO: tl.constexpr,
+    ):
+        # Compute causal m_block range for this n_block
+        (
+            m_block_min,
+            m_block_max,
+            m_block_window_min,
+            m_block_window_max,
+            m_block_sink_min,
+            m_block_sink_max,
+        ) = block_info.get_m_block_min_max(
+            seqlen_q=config.actual_seqlen_q,
+            seqlen_k=config.actual_seqlen_k,
+            n_block=config.n_block,
+            split_idx=split_idx,
+            num_splits=num_splits,
+            window_size_sink=config.window_size_sink,
+            window_size_left=config.window_size_left,
+            window_size_right=config.window_size_right,
+            window_size_dist=config.window_size_dist,
+            TILE_N=config.TILE_N,
+            TILE_M=config.TILE_M,
+            IS_CAUSAL=IS_CAUSAL,
+            IS_LOCAL=IS_LOCAL,
+            IS_SPLIT_QO=IS_SPLIT_QO,
+        )
+        m_block_min_no_mask = block_info.get_m_block_min_causal_local_mask(
+            seqlen_q=config.actual_seqlen_q,
+            seqlen_k=config.actual_seqlen_k,
+            n_block=config.n_block,
+            m_block_min=m_block_min,
+            window_size_right=0,
+            window_size_dist=0,
+            TILE_N=config.TILE_N,
+            TILE_M=config.TILE_M,
+            IS_CAUSAL=IS_CAUSAL or IS_LOCAL,
+            IS_LOCAL=False,
+        )
+        # Clamp to split's range so the no-mask loop stays within bounds
+        if IS_SPLIT_QO:
+            m_block_min_no_mask = tl.minimum(m_block_min_no_mask, m_block_max)
+        if IS_LOCAL:
+            # Compute local m_block range for this n_block
+            m_block_min_no_mask = m_block_max
+            m_block_window_min = tl.maximum(m_block_window_min, m_block_min_no_mask)
+            m_block_window_min_no_mask = block_info.get_m_block_min_causal_local_mask(
+                seqlen_q=config.actual_seqlen_q,
+                seqlen_k=config.actual_seqlen_k,
+                n_block=config.n_block,
+                m_block_min=m_block_window_min,
+                window_size_right=config.window_size_right,
+                window_size_dist=config.window_size_dist,
+                TILE_N=config.TILE_N,
+                TILE_M=config.TILE_M,
+                IS_CAUSAL=False,
+                IS_LOCAL=True,
+            )
+            m_block_window_min_no_mask = tl.maximum(
+                m_block_window_min_no_mask, m_block_window_min
+            )
+            m_block_window_max_no_mask = block_info.get_m_block_max_before_local_mask(
+                seqlen_q=config.actual_seqlen_q,
+                seqlen_k=config.actual_seqlen_k,
+                n_block=config.n_block,
+                m_block_max=m_block_window_max,
+                window_size_left=config.window_size_left,
+                window_size_right=config.window_size_right,
+                window_size_dist=config.window_size_dist,
+                TILE_N=config.TILE_N,
+                TILE_M=config.TILE_M,
+                IS_LOCAL=True,
+            )
+            m_block_window_max_no_mask = tl.maximum(
+                m_block_window_max_no_mask, m_block_window_min_no_mask
+            )
+            m_block_window_min_no_mask = tl.maximum(
+                tl.minimum(m_block_window_min_no_mask, m_block_window_max),
+                m_block_window_min,
+            )
+            m_block_window_max_no_mask = tl.maximum(
+                tl.minimum(m_block_window_max_no_mask, m_block_window_max),
+                m_block_window_min_no_mask,
+            )
+        else:
+            m_block_window_min_no_mask = 0
+            m_block_window_max_no_mask = 0
+            m_block_sink_min = 0
+            m_block_sink_max = 0
+        return AttnBwdBlockScheduler(
+            m_block_min,
+            m_block_max,
+            m_block_min_no_mask,
+            m_block_window_min,
+            m_block_window_max,
+            m_block_window_min_no_mask,
+            m_block_window_max_no_mask,
+            m_block_sink_min,
+            m_block_sink_max,
+        )
+@aggregate
+class AttnDecBlockScheduler:
+    n_block_min: tl.tensor
+    n_block_max: tl.tensor
+    n_block_max_no_mask: tl.tensor
+    n_block_window_min: tl.tensor
+    n_block_window_max: tl.tensor
+    n_block_window_min_no_mask: tl.tensor
+    n_block_window_max_no_mask: tl.tensor
+    n_block_sink_min: tl.tensor
+    n_block_sink_max: tl.tensor
+    @constexpr_function
+    def __init__(
+        self,
+        n_block_min,
+        n_block_max,
+        n_block_max_no_mask,
+        n_block_window_min,
+        n_block_window_max,
+        n_block_window_min_no_mask,
+        n_block_window_max_no_mask,
+        n_block_sink_min,
+        n_block_sink_max,
+    ):
+        self.n_block_min = n_block_min
+        self.n_block_max = n_block_max
+        self.n_block_max_no_mask = n_block_max_no_mask
+        self.n_block_window_min = n_block_window_min
+        self.n_block_window_max = n_block_window_max
+        self.n_block_window_min_no_mask = n_block_window_min_no_mask
+        self.n_block_window_max_no_mask = n_block_window_max_no_mask
+        self.n_block_sink_min = n_block_sink_min
+        self.n_block_sink_max = n_block_sink_max
+    @triton.jit
+    def is_empty(self):
+        return (
+            (self.n_block_max <= self.n_block_min)
+            & (self.n_block_window_max <= self.n_block_window_min)
+            & (self.n_block_sink_max <= self.n_block_sink_min)
+        )
+    @staticmethod
+    @triton.jit
+    def create(
+        config: AttnDecConfig,
+        split_idx,
+        num_splits,
+        IS_LOCAL: tl.constexpr,
+    ):
+        # Compute non-causal n_block range for this m_block
+        (
+            n_block_min,
+            n_block_max,
+            n_block_window_min,
+            n_block_window_max,
+            n_block_sink_min,
+            n_block_sink_max,
+        ) = block_info.get_n_block_min_max(
+            seqlen_q=1,
+            seqlen_k=config.actual_seqlen_k,
+            m_block=0,
+            split_idx=split_idx,
+            num_splits=num_splits,
+            window_size_sink=config.window_size_sink,
+            window_size_left=config.window_size_left,
+            window_size_right=config.window_size_right,
+            window_size_dist=config.window_size_dist,
+            TILE_N=config.TILE_N,
+            TILE_M=config.TILE_M,
+            IS_CAUSAL=False,
+            IS_LOCAL=IS_LOCAL,
+            IS_SPLIT_KV=True,
+            QHEAD_PER_KVHEAD_PACKGQA=config.QHEAD_PER_KVHEAD_PACKGQA,
+        )
+        n_block_max_no_mask = block_info.get_n_block_min_causal_local_mask(
+            seqlen_q=1,
+            seqlen_k=config.actual_seqlen_k,
+            m_block=0,
+            n_block_min=n_block_min,
+            window_size_right=0,
+            window_size_dist=0,
+            TILE_N=config.TILE_N,
+            TILE_M=config.TILE_M,
+            IS_LOCAL=False,
+            QHEAD_PER_KVHEAD_PACKGQA=config.QHEAD_PER_KVHEAD_PACKGQA,
+        )
+        # Clamp to split's range so the no-mask loop stays within bounds
+        n_block_max_no_mask = tl.minimum(n_block_max_no_mask, n_block_max)
+        if IS_LOCAL:
+            # Compute local n_block range for this m_block
+            n_block_max_no_mask = tl.where(
+                split_idx >= num_splits - 1,
+                n_block_min,
+                n_block_max_no_mask,
+            )
+            n_block_window_min = tl.maximum(n_block_window_min, n_block_min)
+            n_block_window_max = tl.minimum(n_block_window_max, n_block_max_no_mask)
+            n_block_window_max_no_mask = block_info.get_n_block_min_causal_local_mask(
+                seqlen_q=1,
+                seqlen_k=config.actual_seqlen_k,
+                m_block=0,
+                n_block_min=n_block_window_min,
+                window_size_right=config.window_size_right,
+                window_size_dist=config.window_size_dist,
+                TILE_N=config.TILE_N,
+                TILE_M=config.TILE_M,
+                IS_LOCAL=True,
+                QHEAD_PER_KVHEAD_PACKGQA=config.QHEAD_PER_KVHEAD_PACKGQA,
+            )
+            n_block_window_min_no_mask = block_info.get_n_block_min_before_local_mask(
+                seqlen_q=1,
+                seqlen_k=config.actual_seqlen_k,
+                m_block=0,
+                n_block_min=n_block_window_min,
+                window_size_left=config.window_size_left,
+                window_size_right=config.window_size_right,
+                window_size_dist=config.window_size_dist,
+                TILE_N=config.TILE_N,
+                TILE_M=config.TILE_M,
+                IS_LOCAL=True,
+                QHEAD_PER_KVHEAD_PACKGQA=config.QHEAD_PER_KVHEAD_PACKGQA,
+            )
+            n_block_window_min_no_mask = tl.minimum(
+                n_block_window_min_no_mask, n_block_window_max_no_mask
+            )
+            # Clamp window no-mask boundaries to the window's range
+            n_block_window_max_no_mask = tl.maximum(
+                tl.minimum(n_block_window_max_no_mask, n_block_window_max),
+                n_block_window_min,
+            )
+            n_block_window_min_no_mask = tl.maximum(
+                tl.minimum(n_block_window_min_no_mask, n_block_window_max),
+                n_block_window_min,
+            )
+        else:
+            n_block_window_min = 0
+            n_block_window_max = 0
+            n_block_window_min_no_mask = 0
+            n_block_window_max_no_mask = 0
+        return AttnDecBlockScheduler(
+            n_block_min,
+            n_block_max,
+            n_block_max_no_mask,
+            n_block_window_min,
+            n_block_window_max,
+            n_block_window_min_no_mask,
+            n_block_window_max_no_mask,
+            n_block_sink_min,
+            n_block_sink_max,
+        )
+@aggregate
+class AttnFwdPointerScheduler:
+    q_base: tl.tensor
+    k_base: tl.tensor
+    v_base: tl.tensor
+    a_base: tl.tensor
+    d_base: tl.tensor
+    out_base: tl.tensor
+    lse_base: tl.tensor
+    head_idx: tl.tensor
+    stride_qh: tl.tensor
+    stride_qm: tl.tensor
+    stride_kn: tl.tensor
+    stride_vn: tl.tensor
+    stride_ah: tl.tensor
+    stride_oh: tl.tensor
+    stride_om: tl.tensor
+    stride_lh: tl.tensor
+    @constexpr_function
+    def __init__(
+        self,
+        q_base,
+        k_base,
+        v_base,
+        a_base,
+        d_base,
+        out_base,
+        lse_base,
+        head_idx,
+        stride_qh,
+        stride_qm,
+        stride_kn,
+        stride_vn,
+        stride_ah,
+        stride_oh,
+        stride_om,
+        stride_lh,
+    ):
+        self.q_base = q_base
+        self.k_base = k_base
+        self.v_base = v_base
+        self.a_base = a_base
+        self.d_base = d_base
+        self.out_base = out_base
+        self.lse_base = lse_base
+        self.head_idx = head_idx
+        self.stride_qh = stride_qh
+        self.stride_qm = stride_qm
+        self.stride_kn = stride_kn
+        self.stride_vn = stride_vn
+        self.stride_ah = stride_ah
+        self.stride_oh = stride_oh
+        self.stride_om = stride_om
+        self.stride_lh = stride_lh
+    @staticmethod
+    @triton.jit
+    def create(
+        config: AttnFwdConfig,
+        Q=None,
+        K=None,
+        V=None,
+        A=None,
+        D=None,
+        Out=None,
+        Lse=None,
+        batch_idx=0,
+        head_idx=0,
+        head_kv_idx=0,
+        split_idx=0,
+        stride_qb=0,
+        stride_qh=0,
+        stride_qm=0,
+        stride_kb=0,
+        stride_kh=0,
+        stride_kn=0,
+        stride_vb=0,
+        stride_vh=0,
+        stride_vn=0,
+        stride_ab=0,
+        stride_ah=0,
+        stride_am=0,
+        stride_db=0,
+        stride_dh=0,
+        stride_dn=0,
+        stride_ob=0,
+        stride_oh=0,
+        stride_om=0,
+        stride_os=0,
+        stride_lb=0,
+        stride_lh=0,
+        stride_ls=0,
+        IS_SPLIT_KV: tl.constexpr = False,
+        IS_GATED: tl.constexpr = False,
+        HAS_CU_SEQLENS_Q: tl.constexpr = False,
+        HAS_CU_SEQLENS_K: tl.constexpr = False,
+    ):
+        # Initialize base pointers
+        q_base = seqlen_info.offset_batch_Q(
+            Q + head_idx * stride_qh if not config.PACK_GQA else Q,
+            batch_idx,
+            config.offset_q,
+            config.padded_offset_q,
+            stride_qb,
+            stride_qm,
+            HAS_CU_SEQLENS_Q,
+            USE_PADDED=False,
+        )
+        k_base = seqlen_info.offset_batch_K(
+            K + head_kv_idx * stride_kh,
+            batch_idx,
+            config.offset_k,
+            config.padded_offset_k,
+            stride_kb,
+            stride_kn,
+            HAS_CU_SEQLENS_K,
+            USE_PADDED=False,
+        )
+        v_base = seqlen_info.offset_batch_K(
+            V + head_kv_idx * stride_vh,
+            batch_idx,
+            config.offset_k,
+            config.padded_offset_k,
+            stride_vb,
+            stride_vn,
+            HAS_CU_SEQLENS_K,
+            USE_PADDED=False,
+        )
+        out_base = seqlen_info.offset_batch_Q(
+            Out + head_idx * stride_oh if not config.PACK_GQA else Out,
+            batch_idx,
+            config.offset_q,
+            config.padded_offset_q,
+            stride_ob,
+            stride_om,
+            HAS_CU_SEQLENS_Q,
+            USE_PADDED=False,
+        )
+        lse_base = seqlen_info.offset_batch_Q(
+            Lse + head_idx * stride_lh if not config.PACK_GQA else Lse,
+            batch_idx,
+            config.offset_q,
+            config.padded_offset_q,
+            stride_lb,
+            1,
+            HAS_CU_SEQLENS_Q,
+            USE_PADDED=False,
+        )
+        if IS_GATED:
+            a_base = seqlen_info.offset_batch_Q(
+                A + head_idx * stride_ah if not config.PACK_GQA else A,
+                batch_idx,
+                config.offset_q,
+                config.padded_offset_q,
+                stride_ab,
+                stride_am,
+                HAS_CU_SEQLENS_Q,
+                USE_PADDED=False,
+            )
+            d_base = seqlen_info.offset_batch_K(
+                D + head_kv_idx * stride_dh,
+                batch_idx,
+                config.offset_k,
+                config.padded_offset_k,
+                stride_db,
+                stride_dn,
+                HAS_CU_SEQLENS_K,
+                USE_PADDED=False,
+            )
+        else:
+            a_base = tl.full((), 0, tl.int64)
+            d_base = tl.full((), 0, tl.int64)
+            stride_ah = tl.full((), 0, tl.int64)
+        # For split KV, offset output and LSE base pointers by split_idx
+        if IS_SPLIT_KV:
+            out_base += split_idx * stride_os
+            lse_base += split_idx * stride_ls
+        # TODO: remove once Triton aggregate accepts host scalar strides or callers materialize strides
+        return AttnFwdPointerScheduler(
+            q_base,
+            k_base,
+            v_base,
+            a_base,
+            d_base,
+            out_base,
+            lse_base,
+            head_idx,
+            stride_qh + tl.full((), 0, tl.int64),
+            stride_qm + tl.full((), 0, tl.int64),
+            stride_kn + tl.full((), 0, tl.int64),
+            stride_vn + tl.full((), 0, tl.int64),
+            stride_ah + tl.full((), 0, tl.int64),
+            stride_oh + tl.full((), 0, tl.int64),
+            stride_om + tl.full((), 0, tl.int64),
+            stride_lh + tl.full((), 0, tl.int64),
+        )
+    @triton.jit
+    def make_q_ptrs(self, config: AttnFwdConfig):
+        if config.PACK_GQA:
+            return seqlen_info.make_pack_gqa_ptrs(
+                self.q_base,
+                config.m_block,
+                self.head_idx,
+                self.stride_qh,
+                self.stride_qm,
+                TILE_M=config.TILE_M,
+                TILE_K=config.TILE_K,
+                QHEAD_PER_KVHEAD_PACKGQA=config.QHEAD_PER_KVHEAD_PACKGQA,
+            )
+        else:
+            return tl.make_tensor_descriptor(
+                self.q_base,
+                shape=[config.actual_seqlen_q, config.head_dim],
+                strides=[self.stride_qm, 1],
+                block_shape=[config.TILE_M, config.TILE_K],
+            )
+    @triton.jit
+    def make_k_ptrs(self, config: AttnFwdConfig):
+        return tl.make_tensor_descriptor(
+            self.k_base,
+            shape=[config.actual_seqlen_k, config.head_dim],
+            strides=[self.stride_kn, 1],
+            block_shape=[config.TILE_N, config.TILE_K],
+        )
+    @triton.jit
+    def make_v_ptrs(self, config: AttnFwdConfig):
+        return tl.make_tensor_descriptor(
+            self.v_base,
+            shape=[config.actual_seqlen_k, config.head_dim],
+            strides=[self.stride_vn, 1],
+            block_shape=[config.TILE_N, config.TILE_K],
+        )
+    @triton.jit
+    def make_a_ptrs(self, config: AttnFwdConfig):
+        if config.PACK_GQA:
+            return seqlen_info.make_pack_gqa_ptrs(
+                self.a_base,
+                config.m_block,
+                self.head_idx,
+                self.stride_ah,
+                1,
+                TILE_M=config.TILE_M,
+                TILE_K=1,
+                QHEAD_PER_KVHEAD_PACKGQA=config.QHEAD_PER_KVHEAD_PACKGQA,
+            )
+        else:
+            return tl.make_tensor_descriptor(
+                self.a_base,
+                shape=[config.actual_seqlen_q],
+                strides=[1],
+                block_shape=[config.TILE_M],
+            )
+    @triton.jit
+    def make_d_ptrs(self, config: AttnFwdConfig):
+        return tl.make_tensor_descriptor(
+            self.d_base,
+            shape=[config.actual_seqlen_k],
+            strides=[1],
+            block_shape=[config.TILE_N],
+        )
+    @triton.jit
+    def make_out_ptrs(self, config: AttnFwdConfig):
+        if config.PACK_GQA:
+            return seqlen_info.make_pack_gqa_ptrs(
+                self.out_base,
+                config.m_block,
+                self.head_idx,
+                self.stride_oh,
+                self.stride_om,
+                TILE_M=config.TILE_M,
+                TILE_K=config.TILE_K,
+                QHEAD_PER_KVHEAD_PACKGQA=config.QHEAD_PER_KVHEAD_PACKGQA,
+            )
+        else:
+            return tl.make_tensor_descriptor(
+                self.out_base,
+                shape=[config.actual_seqlen_q, config.head_dim],
+                strides=[self.stride_om, 1],
+                block_shape=[config.TILE_M, config.TILE_K],
+            )
+    @triton.jit
+    def make_lse_ptrs(self, config: AttnFwdConfig):
+        if config.PACK_GQA:
+            return seqlen_info.make_pack_gqa_ptrs(
+                self.lse_base,
+                config.m_block,
+                self.head_idx,
+                self.stride_lh,
+                1,
+                TILE_M=config.TILE_M,
+                TILE_K=1,
+                QHEAD_PER_KVHEAD_PACKGQA=config.QHEAD_PER_KVHEAD_PACKGQA,
+            )
+        else:
+            return tl.make_tensor_descriptor(
+                self.lse_base,
+                shape=[config.actual_seqlen_q],
+                strides=[1],
+                block_shape=[config.TILE_M],
+            )
+    @triton.jit
+    def load_q(self, config: AttnFwdConfig, q_ptrs):
+        if config.PACK_GQA:
+            offs_m = config.get_offs_m()
+            offs_k = config.get_offs_k()
+            return tl.load(
+                q_ptrs,
+                mask=(
+                    (offs_m // config.QHEAD_PER_KVHEAD_PACKGQA) < config.actual_seqlen_q
+                )[:, None]
+                & (offs_k < config.head_dim)[None, :],
+                other=0.0,
+                cache_modifier=".ca",
+            )
+        else:
+            return q_ptrs.load([config.m_block * config.TILE_M, 0])
+    @triton.jit
+    def load_k(self, config: AttnFwdConfig, k_ptrs, n_block):
+        return k_ptrs.load([n_block * config.TILE_N, 0])
+    @triton.jit
+    def load_v(self, config: AttnFwdConfig, v_ptrs, n_block):
+        return v_ptrs.load([n_block * config.TILE_N, 0])
+    @triton.jit
+    def load_a(self, config: AttnFwdConfig, a_ptrs):
+        if config.PACK_GQA:
+            offs_m = config.get_offs_m()
+            return tl.load(
+                a_ptrs,
+                mask=(offs_m // config.QHEAD_PER_KVHEAD_PACKGQA)
+                < config.actual_seqlen_q,
+                other=0.0,
+            ).to(tl.float32)
+        else:
+            return a_ptrs.load([config.m_block * config.TILE_M]).to(tl.float32)
+    @triton.jit
+    def load_d(self, config: AttnFwdConfig, d_ptrs, n_block):
+        return d_ptrs.load([n_block * config.TILE_N]).to(tl.float32)
+    @triton.jit
+    def store_out(
+        self, config: AttnFwdConfig, out_ptrs, o_tile, IS_SPLIT_KV: tl.constexpr = False
+    ):
+        if not IS_SPLIT_KV:
+            if config.PACK_GQA:
+                o_tile = o_tile.to(out_ptrs.dtype.element_ty)
+            else:
+                o_tile = o_tile.to(out_ptrs.dtype)
+        if config.PACK_GQA:
+            offs_m = config.get_offs_m()
+            offs_k = config.get_offs_k()
+            tl.store(
+                out_ptrs,
+                o_tile,
+                mask=(
+                    (offs_m // config.QHEAD_PER_KVHEAD_PACKGQA) < config.actual_seqlen_q
+                )[:, None]
+                & (offs_k < config.head_dim)[None, :],
+                cache_modifier=".wb",
+            )
+        else:
+            out_ptrs.store([config.m_block * config.TILE_M, 0], o_tile)
+    @triton.jit
+    def store_lse(self, config: AttnFwdConfig, lse_ptrs, lse_tile):
+        if config.PACK_GQA:
+            offs_m = config.get_offs_m()
+            tl.store(
+                lse_ptrs,
+                lse_tile,
+                mask=(
+                    (offs_m // config.QHEAD_PER_KVHEAD_PACKGQA) < config.actual_seqlen_q
+                ),
+                cache_modifier=".wb",
+            )
+        else:
+            lse_ptrs.store([config.m_block * config.TILE_M], lse_tile)
+    @triton.jit
+    def store_empty(
+        self,
+        config: AttnFwdConfig,
+        out_ptrs,
+        lse_ptrs,
+        IS_SPLIT_KV: tl.constexpr = False,
+    ):
+        lse_tile = tl.full((config.TILE_M,), float("-inf"), dtype=tl.float32)
+        self.store_lse(config, lse_ptrs, lse_tile)
+        o_tile = tl.zeros((config.TILE_M, config.TILE_K), dtype=tl.float32)
+        self.store_out(config, out_ptrs, o_tile, IS_SPLIT_KV=IS_SPLIT_KV)
+@aggregate
+class AttnBwdPointerScheduler:
+    q_base: tl.tensor
+    k_base: tl.tensor
+    v_base: tl.tensor
+    a_base: tl.tensor
+    d_base: tl.tensor
+    do_base: tl.tensor
+    lse_base: tl.tensor
+    dpsum_base: tl.tensor
+    dq_accum_base: tl.tensor
+    dk_base: tl.tensor
+    dv_base: tl.tensor
+    da_base: tl.tensor
+    dd_base: tl.tensor
+    stride_qm: tl.tensor
+    stride_kn: tl.tensor
+    stride_vn: tl.tensor
+    stride_dom: tl.tensor
+    stride_dqam: tl.tensor
+    stride_dkn: tl.tensor
+    stride_dvn: tl.tensor
+    @constexpr_function
+    def __init__(
+        self,
+        q_base,
+        k_base,
+        v_base,
+        a_base,
+        d_base,
+        do_base,
+        lse_base,
+        dpsum_base,
+        dq_accum_base,
+        dk_base,
+        dv_base,
+        da_base,
+        dd_base,
+        stride_qm,
+        stride_kn,
+        stride_vn,
+        stride_dom,
+        stride_dqam,
+        stride_dkn,
+        stride_dvn,
+    ):
+        self.q_base = q_base
+        self.k_base = k_base
+        self.v_base = v_base
+        self.a_base = a_base
+        self.d_base = d_base
+        self.do_base = do_base
+        self.lse_base = lse_base
+        self.dpsum_base = dpsum_base
+        self.dq_accum_base = dq_accum_base
+        self.dk_base = dk_base
+        self.dv_base = dv_base
+        self.da_base = da_base
+        self.dd_base = dd_base
+        self.stride_qm = stride_qm
+        self.stride_kn = stride_kn
+        self.stride_vn = stride_vn
+        self.stride_dom = stride_dom
+        self.stride_dqam = stride_dqam
+        self.stride_dkn = stride_dkn
+        self.stride_dvn = stride_dvn
+    @staticmethod
+    @triton.jit
+    def create(
+        config: AttnBwdConfig,
+        Q=None,
+        K=None,
+        V=None,
+        A=None,
+        D=None,
+        dO=None,
+        LSELog2=None,
+        dPsum=None,
+        dQaccum=None,
+        dK=None,
+        dV=None,
+        dA=None,
+        dD=None,
+        batch_idx=0,
+        head_idx=0,
+        head_kv_idx=0,
+        split_idx=0,
+        stride_qb=0,
+        stride_qh=0,
+        stride_qm=0,
+        stride_kb=0,
+        stride_kh=0,
+        stride_kn=0,
+        stride_vb=0,
+        stride_vh=0,
+        stride_vn=0,
+        stride_ab=0,
+        stride_ah=0,
+        stride_db=0,
+        stride_dh=0,
+        stride_dob=0,
+        stride_doh=0,
+        stride_dom=0,
+        stride_lb=0,
+        stride_lh=0,
+        stride_pb=0,
+        stride_ph=0,
+        stride_dqab=0,
+        stride_dqah=0,
+        stride_dqam=0,
+        stride_dkb=0,
+        stride_dkh=0,
+        stride_dkn=0,
+        stride_dks=0,
+        stride_dvb=0,
+        stride_dvh=0,
+        stride_dvn=0,
+        stride_dvs=0,
+        stride_dab=0,
+        stride_dah=0,
+        stride_ddb=0,
+        stride_ddh=0,
+        stride_dds=0,
+        IS_SPLIT_QO: tl.constexpr = False,
+        IS_GATED: tl.constexpr = False,
+        HAS_CU_SEQLENS_Q: tl.constexpr = False,
+        HAS_CU_SEQLENS_K: tl.constexpr = False,
+    ):
+        # Initialize base pointers
+        q_base = seqlen_info.offset_batch_Q(
+            Q + head_idx * stride_qh,
+            batch_idx,
+            config.offset_q,
+            config.padded_offset_q,
+            stride_qb,
+            stride_qm,
+            HAS_CU_SEQLENS_Q,
+            USE_PADDED=False,
+        )
+        k_base = seqlen_info.offset_batch_K(
+            K + head_kv_idx * stride_kh,
+            batch_idx,
+            config.offset_k,
+            config.padded_offset_k,
+            stride_kb,
+            stride_kn,
+            HAS_CU_SEQLENS_K,
+            USE_PADDED=False,
+        )
+        v_base = seqlen_info.offset_batch_K(
+            V + head_kv_idx * stride_vh,
+            batch_idx,
+            config.offset_k,
+            config.padded_offset_k,
+            stride_vb,
+            stride_vn,
+            HAS_CU_SEQLENS_K,
+            USE_PADDED=False,
+        )
+        do_base = seqlen_info.offset_batch_Q(
+            dO + head_idx * stride_doh,
+            batch_idx,
+            config.offset_q,
+            config.padded_offset_q,
+            stride_dob,
+            stride_dom,
+            HAS_CU_SEQLENS_Q,
+            USE_PADDED=False,
+        )
+        lse_base = seqlen_info.offset_batch_Q(
+            LSELog2 + head_idx * stride_lh,
+            batch_idx,
+            config.offset_q,
+            config.padded_offset_q,
+            stride_lb,
+            1,
+            HAS_CU_SEQLENS_Q,
+            USE_PADDED=True,
+        )
+        dpsum_base = seqlen_info.offset_batch_Q(
+            dPsum + head_idx * stride_ph,
+            batch_idx,
+            config.offset_q,
+            config.padded_offset_q,
+            stride_pb,
+            1,
+            HAS_CU_SEQLENS_Q,
+            USE_PADDED=True,
+        )
+        dq_accum_base = seqlen_info.offset_batch_Q(
+            dQaccum + head_idx * stride_dqah,
+            batch_idx,
+            config.offset_q,
+            config.padded_offset_q,
+            stride_dqab,
+            stride_dqam,
+            HAS_CU_SEQLENS_Q,
+            USE_PADDED=True,
+        )
+        dk_base = seqlen_info.offset_batch_K(
+            dK + head_kv_idx * stride_dkh,
+            batch_idx,
+            config.offset_k,
+            config.padded_offset_k,
+            stride_dkb,
+            stride_dkn,
+            HAS_CU_SEQLENS_K,
+            USE_PADDED=False,
+        )
+        dv_base = seqlen_info.offset_batch_K(
+            dV + head_kv_idx * stride_dvh,
+            batch_idx,
+            config.offset_k,
+            config.padded_offset_k,
+            stride_dvb,
+            stride_dvn,
+            HAS_CU_SEQLENS_K,
+            USE_PADDED=False,
+        )
+        if IS_GATED:
+            a_base = seqlen_info.offset_batch_Q(
+                A + head_idx * stride_ah,
+                batch_idx,
+                config.offset_q,
+                config.padded_offset_q,
+                stride_ab,
+                1,
+                HAS_CU_SEQLENS_Q,
+                USE_PADDED=False,
+            )
+            d_base = seqlen_info.offset_batch_K(
+                D + head_kv_idx * stride_dh,
+                batch_idx,
+                config.offset_k,
+                config.padded_offset_k,
+                stride_db,
+                1,
+                HAS_CU_SEQLENS_K,
+                USE_PADDED=False,
+            )
+            da_base = seqlen_info.offset_batch_Q(
+                dA + head_idx * stride_dah,
+                batch_idx,
+                config.offset_q,
+                config.padded_offset_q,
+                stride_dab,
+                1,
+                HAS_CU_SEQLENS_Q,
+                USE_PADDED=True,
+            )
+            dd_base = seqlen_info.offset_batch_K(
+                dD + head_kv_idx * stride_ddh,
+                batch_idx,
+                config.offset_k,
+                config.padded_offset_k,
+                stride_ddb,
+                1,
+                HAS_CU_SEQLENS_K,
+                USE_PADDED=False,
+            )
+        else:
+            a_base = tl.full((), 0, tl.int64)
+            d_base = tl.full((), 0, tl.int64)
+            da_base = tl.full((), 0, tl.int64)
+            dd_base = tl.full((), 0, tl.int64)
+        # For split QO, offset key and value gradients base pointers by split_idx
+        if IS_SPLIT_QO:
+            dk_base += split_idx * stride_dks
+            dv_base += split_idx * stride_dvs
+            if IS_GATED:
+                dd_base += split_idx * stride_dds
+        # TODO: remove once Triton aggregate accepts host scalar strides or callers materialize strides
+        return AttnBwdPointerScheduler(
+            q_base,
+            k_base,
+            v_base,
+            a_base,
+            d_base,
+            do_base,
+            lse_base,
+            dpsum_base,
+            dq_accum_base,
+            dk_base,
+            dv_base,
+            da_base,
+            dd_base,
+            stride_qm + tl.full((), 0, tl.int64),
+            stride_kn + tl.full((), 0, tl.int64),
+            stride_vn + tl.full((), 0, tl.int64),
+            stride_dom + tl.full((), 0, tl.int64),
+            stride_dqam + tl.full((), 0, tl.int64),
+            stride_dkn + tl.full((), 0, tl.int64),
+            stride_dvn + tl.full((), 0, tl.int64),
+        )
+    @triton.jit
+    def make_q_ptrs(self, config: AttnBwdConfig):
+        return tl.make_tensor_descriptor(
+            self.q_base,
+            shape=[config.actual_seqlen_q, config.head_dim],
+            strides=[self.stride_qm, 1],
+            block_shape=[config.TILE_M, config.TILE_K],
+        )
+    @triton.jit
+    def make_k_ptrs(self, config: AttnBwdConfig):
+        return tl.make_tensor_descriptor(
+            self.k_base,
+            shape=[config.actual_seqlen_k, config.head_dim],
+            strides=[self.stride_kn, 1],
+            block_shape=[config.TILE_N, config.TILE_K],
+        )
+    @triton.jit
+    def make_v_ptrs(self, config: AttnBwdConfig):
+        return tl.make_tensor_descriptor(
+            self.v_base,
+            shape=[config.actual_seqlen_k, config.head_dim],
+            strides=[self.stride_vn, 1],
+            block_shape=[config.TILE_N, config.TILE_K],
+        )
+    @triton.jit
+    def make_a_ptrs(self, config: AttnBwdConfig):
+        return tl.make_tensor_descriptor(
+            self.a_base,
+            shape=[config.actual_seqlen_q],
+            strides=[1],
+            block_shape=[config.TILE_M],
+        )
+    @triton.jit
+    def make_d_ptrs(self, config: AttnBwdConfig):
+        return tl.make_tensor_descriptor(
+            self.d_base,
+            shape=[config.actual_seqlen_k],
+            strides=[1],
+            block_shape=[config.TILE_N],
+        )
+    @triton.jit
+    def make_do_ptrs(self, config: AttnBwdConfig):
+        return tl.make_tensor_descriptor(
+            self.do_base,
+            shape=[config.actual_seqlen_q, config.head_dim],
+            strides=[self.stride_dom, 1],
+            block_shape=[config.TILE_M, config.TILE_K],
+        )
+    @triton.jit
+    def make_lse_ptrs(self, config: AttnBwdConfig):
+        return tl.make_tensor_descriptor(
+            self.lse_base,
+            shape=[config.actual_seqlen_q],
+            strides=[1],
+            block_shape=[config.TILE_M],
+        )
+    @triton.jit
+    def make_dpsum_ptrs(self, config: AttnBwdConfig):
+        return tl.make_tensor_descriptor(
+            self.dpsum_base,
+            shape=[config.actual_seqlen_q],
+            strides=[1],
+            block_shape=[config.TILE_M],
+        )
+    @triton.jit
+    def make_dq_accum_ptrs(self, config: AttnBwdConfig):
+        return tl.make_tensor_descriptor(
+            self.dq_accum_base,
+            shape=[config.actual_seqlen_q, config.head_dim],
+            strides=[self.stride_dqam, 1],
+            block_shape=[config.TILE_M, config.TILE_K],
+        )
+    @triton.jit
+    def make_dk_ptrs(self, config: AttnBwdConfig):
+        return tl.make_tensor_descriptor(
+            self.dk_base,
+            shape=[config.actual_seqlen_k, config.head_dim],
+            strides=[self.stride_dkn, 1],
+            block_shape=[config.TILE_N, config.TILE_K],
+        )
+    @triton.jit
+    def make_dv_ptrs(self, config: AttnBwdConfig):
+        return tl.make_tensor_descriptor(
+            self.dv_base,
+            shape=[config.actual_seqlen_k, config.head_dim],
+            strides=[self.stride_dvn, 1],
+            block_shape=[config.TILE_N, config.TILE_K],
+        )
+    @triton.jit
+    def make_da_ptrs(self, config: AttnBwdConfig):
+        return tl.make_tensor_descriptor(
+            self.da_base,
+            shape=[config.actual_seqlen_q],
+            strides=[1],
+            block_shape=[config.TILE_M],
+        )
+    @triton.jit
+    def make_dd_ptrs(self, config: AttnBwdConfig):
+        return tl.make_tensor_descriptor(
+            self.dd_base,
+            shape=[config.actual_seqlen_k],
+            strides=[1],
+            block_shape=[config.TILE_N],
+        )
+    @triton.jit
+    def load_q(self, config: AttnBwdConfig, q_ptrs, m_block):
+        return (q_ptrs.load([m_block * config.TILE_M, 0]) * config.query_scale).to(
+            self.do_base.dtype.element_ty
+        )
+    @triton.jit
+    def load_k(self, config: AttnBwdConfig, k_ptrs):
+        return (k_ptrs.load([config.n_block * config.TILE_N, 0]) * config.key_scale).to(
+            self.do_base.dtype.element_ty
+        )
+    @triton.jit
+    def load_v(self, config: AttnBwdConfig, v_ptrs):
+        return (
+            v_ptrs.load([config.n_block * config.TILE_N, 0]) * config.value_scale
+        ).to(self.do_base.dtype.element_ty)
+    @triton.jit
+    def load_a(self, config: AttnBwdConfig, a_ptrs, m_block):
+        return a_ptrs.load([m_block * config.TILE_M]).to(tl.float32)
+    @triton.jit
+    def load_d(self, config: AttnBwdConfig, d_ptrs):
+        return d_ptrs.load([config.n_block * config.TILE_N]).to(tl.float32)
+    @triton.jit
+    def load_do(self, config: AttnBwdConfig, do_ptrs, m_block):
+        return do_ptrs.load([m_block * config.TILE_M, 0])
+    @triton.jit
+    def load_lse(self, config: AttnBwdConfig, lse_ptrs, m_block):
+        return lse_ptrs.load([m_block * config.TILE_M])
+    @triton.jit
+    def load_dpsum(self, config: AttnBwdConfig, dpsum_ptrs, m_block):
+        return dpsum_ptrs.load([m_block * config.TILE_M])
+    @triton.jit
+    def store_dq(self, config: AttnBwdConfig, dq_ptrs, m_block, dq_tile):
+        dq_ptrs.atomic_add([m_block * config.TILE_M, 0], dq_tile)
+    @triton.jit
+    def store_dk(self, config: AttnBwdConfig, dk_ptrs, dk_tile):
+        if config.QHEAD_PER_KVHEAD > 1:
+            dk_ptrs.atomic_add([config.n_block * config.TILE_N, 0], dk_tile)
+        else:
+            dk_ptrs.store([config.n_block * config.TILE_N, 0], dk_tile)
+    @triton.jit
+    def store_dv(self, config: AttnBwdConfig, dv_ptrs, dv_tile):
+        if config.QHEAD_PER_KVHEAD > 1:
+            dv_ptrs.atomic_add([config.n_block * config.TILE_N, 0], dv_tile)
+        else:
+            dv_ptrs.store([config.n_block * config.TILE_N, 0], dv_tile)
+    @triton.jit
+    def store_da(self, config: AttnBwdConfig, da_ptrs, m_block, da_tile):
+        da_ptrs.atomic_add([m_block * config.TILE_M], da_tile)
+    @triton.jit
+    def store_dd(self, config: AttnBwdConfig, dd_ptrs, dd_tile):
+        if config.QHEAD_PER_KVHEAD > 1:
+            dd_ptrs.atomic_add([config.n_block * config.TILE_N], dd_tile)
+        else:
+            dd_ptrs.store([config.n_block * config.TILE_N], dd_tile)
+@aggregate
+class AttnDecPointerScheduler:
+    q_base: tl.tensor
+    k_base: tl.tensor
+    v_base: tl.tensor
+    a_base: tl.tensor
+    d_base: tl.tensor
+    out_base: tl.tensor
+    lse_base: tl.tensor
+    stride_qh: tl.tensor
+    stride_kn: tl.tensor
+    stride_vn: tl.tensor
+    stride_oh: tl.tensor
+    @constexpr_function
+    def __init__(
+        self,
+        q_base,
+        k_base,
+        v_base,
+        a_base,
+        d_base,
+        out_base,
+        lse_base,
+        stride_qh,
+        stride_kn,
+        stride_vn,
+        stride_oh,
+    ):
+        self.q_base = q_base
+        self.k_base = k_base
+        self.v_base = v_base
+        self.a_base = a_base
+        self.d_base = d_base
+        self.out_base = out_base
+        self.lse_base = lse_base
+        self.stride_qh = stride_qh
+        self.stride_kn = stride_kn
+        self.stride_vn = stride_vn
+        self.stride_oh = stride_oh
+    @staticmethod
+    @triton.jit
+    def create(
+        config: AttnDecConfig,
+        Q=None,
+        K=None,
+        V=None,
+        A=None,
+        D=None,
+        Out=None,
+        Lse=None,
+        batch_idx=0,
+        head_idx=0,
+        head_kv_idx=0,
+        split_idx=0,
+        stride_qb=0,
+        stride_qh=0,
+        stride_qm=0,
+        stride_kb=0,
+        stride_kh=0,
+        stride_kn=0,
+        stride_vb=0,
+        stride_vh=0,
+        stride_vn=0,
+        stride_ab=0,
+        stride_ah=0,
+        stride_db=0,
+        stride_dh=0,
+        stride_ob=0,
+        stride_oh=0,
+        stride_om=0,
+        stride_os=0,
+        stride_lb=0,
+        stride_lh=0,
+        stride_lm=0,
+        stride_ls=0,
+        IS_GATED: tl.constexpr = False,
+        HAS_CU_SEQLENS_Q: tl.constexpr = False,
+        HAS_CU_SEQLENS_K: tl.constexpr = False,
+    ):
+        # Initialize base pointers
+        q_base = seqlen_info.offset_batch_Q(
+            Q + head_idx * config.QHEAD_PER_KVHEAD_PACKGQA * stride_qh,
+            batch_idx,
+            config.offset_q,
+            config.padded_offset_q,
+            stride_qb,
+            stride_qm,
+            HAS_CU_SEQLENS_Q,
+            USE_PADDED=False,
+        )
+        k_base = seqlen_info.offset_batch_K(
+            K + head_kv_idx * stride_kh,
+            batch_idx,
+            config.offset_k,
+            config.padded_offset_k,
+            stride_kb,
+            stride_kn,
+            HAS_CU_SEQLENS_K,
+            USE_PADDED=False,
+        )
+        v_base = seqlen_info.offset_batch_K(
+            V + head_kv_idx * stride_vh,
+            batch_idx,
+            config.offset_k,
+            config.padded_offset_k,
+            stride_vb,
+            stride_vn,
+            HAS_CU_SEQLENS_K,
+            USE_PADDED=False,
+        )
+        out_base = seqlen_info.offset_batch_Q(
+            Out + head_idx * config.QHEAD_PER_KVHEAD_PACKGQA * stride_oh,
+            batch_idx,
+            config.offset_q,
+            config.padded_offset_q,
+            stride_ob,
+            stride_om,
+            HAS_CU_SEQLENS_Q,
+            USE_PADDED=False,
+        )
+        lse_base = seqlen_info.offset_batch_Q(
+            Lse + head_idx * config.QHEAD_PER_KVHEAD_PACKGQA * stride_lh,
+            batch_idx,
+            config.offset_q,
+            config.padded_offset_q,
+            stride_lb,
+            stride_lm,
+            HAS_CU_SEQLENS_Q,
+            USE_PADDED=False,
+        )
+        if IS_GATED:
+            a_base = seqlen_info.offset_batch_Q(
+                A + head_idx * config.QHEAD_PER_KVHEAD_PACKGQA * stride_ah,
+                batch_idx,
+                config.offset_q,
+                config.padded_offset_q,
+                stride_ab,
+                1,
+                HAS_CU_SEQLENS_Q,
+                USE_PADDED=False,
+            )
+            d_base = seqlen_info.offset_batch_K(
+                D + head_kv_idx * stride_dh,
+                batch_idx,
+                config.offset_k,
+                config.padded_offset_k,
+                stride_db,
+                1,
+                HAS_CU_SEQLENS_K,
+                USE_PADDED=False,
+            )
+        else:
+            a_base = tl.full((), 0, tl.int64)
+            d_base = tl.full((), 0, tl.int64)
+        # For split KV, offset output and LSE base pointers by split_idx
+        out_base += split_idx * stride_os
+        lse_base += split_idx * stride_ls
+        # TODO: remove once Triton aggregate accepts host scalar strides or callers materialize strides
+        return AttnDecPointerScheduler(
+            q_base,
+            k_base,
+            v_base,
+            a_base,
+            d_base,
+            out_base,
+            lse_base,
+            stride_qh + tl.full((), 0, tl.int64),
+            stride_kn + tl.full((), 0, tl.int64),
+            stride_vn + tl.full((), 0, tl.int64),
+            stride_oh + tl.full((), 0, tl.int64),
+        )
+    @triton.jit
+    def make_q_ptrs(self, config: AttnDecConfig):
+        return tl.make_tensor_descriptor(
+            self.q_base,
+            shape=[config.QHEAD_PER_KVHEAD_PACKGQA, config.head_dim],
+            strides=[self.stride_qh, 1],
+            block_shape=[config.TILE_M, config.TILE_K],
+        )
+    @triton.jit
+    def make_k_ptrs(self, config: AttnDecConfig):
+        return tl.make_tensor_descriptor(
+            self.k_base,
+            shape=[config.actual_seqlen_k, config.head_dim],
+            strides=[self.stride_kn, 1],
+            block_shape=[config.TILE_N, config.TILE_K],
+        )
+    @triton.jit
+    def make_v_ptrs(self, config: AttnDecConfig):
+        return tl.make_tensor_descriptor(
+            self.v_base,
+            shape=[config.actual_seqlen_k, config.head_dim],
+            strides=[self.stride_vn, 1],
+            block_shape=[config.TILE_N, config.TILE_K],
+        )
+    @triton.jit
+    def make_a_ptrs(self, config: AttnDecConfig):
+        return tl.make_tensor_descriptor(
+            self.a_base,
+            shape=[config.QHEAD_PER_KVHEAD_PACKGQA],
+            strides=[1],
+            block_shape=[config.TILE_M],
+        )
+    @triton.jit
+    def make_d_ptrs(self, config: AttnDecConfig):
+        return tl.make_tensor_descriptor(
+            self.d_base,
+            shape=[config.actual_seqlen_k],
+            strides=[1],
+            block_shape=[config.TILE_N],
+        )
+    @triton.jit
+    def make_out_ptrs(self, config: AttnDecConfig):
+        return tl.make_tensor_descriptor(
+            self.out_base,
+            shape=[config.QHEAD_PER_KVHEAD_PACKGQA, config.head_dim],
+            strides=[self.stride_oh, 1],
+            block_shape=[config.TILE_M, config.TILE_K],
+        )
+    @triton.jit
+    def make_lse_ptrs(self, config: AttnDecConfig):
+        return tl.make_tensor_descriptor(
+            self.lse_base,
+            shape=[config.QHEAD_PER_KVHEAD_PACKGQA],
+            strides=[1],
+            block_shape=[config.TILE_M],
+        )
+    @triton.jit
+    def load_q(self, config: AttnDecConfig, q_ptrs):
+        return q_ptrs.load([0, 0])
+    @triton.jit
+    def load_k(self, config: AttnDecConfig, k_ptrs, n_block):
+        return k_ptrs.load([n_block * config.TILE_N, 0])
+    @triton.jit
+    def load_v(self, config: AttnDecConfig, v_ptrs, n_block):
+        return v_ptrs.load([n_block * config.TILE_N, 0])
+    @triton.jit
+    def load_a(self, config: AttnDecConfig, a_ptrs):
+        return a_ptrs.load([0]).to(tl.float32)
+    @triton.jit
+    def load_d(self, config: AttnDecConfig, d_ptrs, n_block):
+        return d_ptrs.load([n_block * config.TILE_N]).to(tl.float32)
+    @triton.jit
+    def store_out(self, config: AttnDecConfig, out_ptrs, o_tile):
+        out_ptrs.store([0, 0], o_tile)
+    @triton.jit
+    def store_lse(self, config: AttnDecConfig, lse_ptrs, lse_tile):
+        lse_ptrs.store([0], lse_tile)
+    @triton.jit
+    def store_empty(self, config: AttnDecConfig, out_ptrs, lse_ptrs, Out):
+        lse_tile = tl.full((config.TILE_M,), float("-inf"), dtype=tl.float32)
+        self.store_lse(config, lse_ptrs, lse_tile)
+        o_tile = tl.zeros((config.TILE_M, config.TILE_K), dtype=Out.dtype.element_ty)
+        self.store_out(config, out_ptrs, o_tile)
+@aggregate
+class AttnMaskScheduler:
+    fixed_block: tl.tensor
+    actual_seqlen_q: tl.tensor
+    actual_seqlen_k: tl.tensor
+    window_size_sink: tl.tensor
+    window_size_left: tl.tensor
+    window_size_right: tl.tensor
+    window_size_dist: tl.tensor
+    TILE_M: tl.constexpr
+    TILE_N: tl.constexpr
+    QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr
+    SWAP_AB: tl.constexpr
+    @constexpr_function
+    def __init__(
+        self,
+        fixed_block,
+        actual_seqlen_q,
+        actual_seqlen_k,
+        window_size_sink,
+        window_size_left,
+        window_size_right,
+        window_size_dist,
+        TILE_M,
+        TILE_N,
+        QHEAD_PER_KVHEAD_PACKGQA,
+        SWAP_AB,
+    ):
+        self.fixed_block = fixed_block
+        self.actual_seqlen_q = actual_seqlen_q
+        self.actual_seqlen_k = actual_seqlen_k
+        self.window_size_sink = window_size_sink
+        self.window_size_left = window_size_left
+        self.window_size_right = window_size_right
+        self.window_size_dist = window_size_dist
+        self.TILE_M = tl.constexpr(TILE_M)
+        self.TILE_N = tl.constexpr(TILE_N)
+        self.QHEAD_PER_KVHEAD_PACKGQA = tl.constexpr(QHEAD_PER_KVHEAD_PACKGQA)
+        self.SWAP_AB = tl.constexpr(SWAP_AB)
+    @staticmethod
+    @triton.jit
+    def create(
+        config,
+        SWAP_AB: tl.constexpr = False,
+    ):
+        if not SWAP_AB:
+            return AttnMaskScheduler(
+                config.m_block,
+                config.actual_seqlen_q,
+                config.actual_seqlen_k,
+                config.window_size_sink,
+                config.window_size_left,
+                config.window_size_right,
+                config.window_size_dist,
+                config.TILE_M,
+                config.TILE_N,
+                config.QHEAD_PER_KVHEAD_PACKGQA,
+                False,
+            )
+        else:
+            return AttnMaskScheduler(
+                config.n_block,
+                config.actual_seqlen_q,
+                config.actual_seqlen_k,
+                config.window_size_sink,
+                config.window_size_left,
+                config.window_size_right,
+                config.window_size_dist,
+                config.TILE_M,
+                config.TILE_N,
+                1,
+                True,
+            )
+    @triton.jit
+    def apply_mask(
+        self,
+        acc_s,
+        iter_block,
+        MASK_CAUSAL: tl.constexpr = False,
+        MASK_LOCAL: tl.constexpr = False,
+        MASK_SINK: tl.constexpr = False,
+    ):
+        if not self.SWAP_AB:
+            m_block = self.fixed_block
+            n_block = iter_block
+        else:
+            m_block = iter_block
+            n_block = self.fixed_block
+        return mask.apply_mask(
+            acc_s=acc_s,
+            m_block=m_block,
+            n_block=n_block,
+            seqlen_q=self.actual_seqlen_q,
+            seqlen_k=self.actual_seqlen_k,
+            window_size_sink=self.window_size_sink,
+            window_size_left=self.window_size_left,
+            window_size_right=self.window_size_right,
+            window_size_dist=self.window_size_dist,
+            MASK_SEQLEN=True,
+            MASK_CAUSAL=MASK_CAUSAL,
+            MASK_LOCAL=MASK_LOCAL,
+            MASK_SINK=MASK_SINK,
+            TILE_M=self.TILE_M,
+            TILE_N=self.TILE_N,
+            QHEAD_PER_KVHEAD_PACKGQA=self.QHEAD_PER_KVHEAD_PACKGQA,
+            SWAP_AB=self.SWAP_AB,
+        )
+@aggregate
+class SoftmaxScheduler:
+    softmax_scale_log2: tl.tensor
+    value_scale: tl.tensor
+    @constexpr_function
+    def __init__(
+        self,
+        softmax_scale_log2,
+        value_scale,
+    ):
+        self.softmax_scale_log2 = softmax_scale_log2
+        self.value_scale = value_scale
+    @staticmethod
+    @triton.jit
+    def create(config):
+        return SoftmaxScheduler(
+            config.softmax_scale_log2,
+            config.value_scale,
+        )
+    @triton.jit
+    def online_softmax(
+        self,
+        acc_s,
+        row_max,
+        row_sum,
+        CHECK_INF: tl.constexpr = False,
+        RESCALE_THRESHOLD: tl.constexpr = 0.0,
+    ):
+        return activations.online_softmax(
+            acc_s=acc_s,
+            row_max=row_max,
+            row_sum=row_sum,
+            scale_log2=self.softmax_scale_log2,
+            CHECK_INF=CHECK_INF,
+            RESCALE_THRESHOLD=RESCALE_THRESHOLD,
+        )
+    @triton.jit
+    def online_sparse_softmax(
+        self,
+        acc_s,
+        row_max,
+        row_sum,
+        softmax_threshold_log2,
+        CHECK_INF: tl.constexpr = False,
+    ):
+        return activations.online_sparse_softmax(
+            acc_s=acc_s,
+            row_max=row_max,
+            row_sum=row_sum,
+            scale_log2=self.softmax_scale_log2,
+            softmax_threshold_log2=softmax_threshold_log2,
+            CHECK_INF=CHECK_INF,
+        )
+    @triton.jit
+    def rescale_o(
+        self,
+        acc_o,
+        row_scale,
+        LAZY_RESCALE: tl.constexpr = False,
+    ):
+        return activations.rescale_o(
+            acc_o=acc_o,
+            row_scale=row_scale,
+            LAZY_RESCALE=LAZY_RESCALE,
+        )
+    @triton.jit
+    def finalize(
+        self,
+        row_max,
+        row_sum,
+        IS_LOG2: tl.constexpr = False,
+        CHECK_NAN: tl.constexpr = True,
+    ):
+        return activations.finalize(
+            row_max=row_max,
+            row_sum=row_sum,
+            scale_log2=self.softmax_scale_log2,
+            final_scale=self.value_scale,
+            IS_LOG2=IS_LOG2,
+            CHECK_NAN=CHECK_NAN,
+        )
+    @triton.jit
+    def log_sigmoid(
+        self,
+        acc_s,
+        FASTMATH: tl.constexpr = False,
+    ):
+        return activations.log_sigmoid(x=acc_s, FASTMATH=FASTMATH)
+    @triton.jit
+    def online_gate(
+        self,
+        a_max,
+        a_min,
+        d_max,
+        d_min,
+        gate_max,
+        gate_threshold_log2,
+    ):
+        return activations.online_gate(
+            a_max=a_max,
+            a_min=a_min,
+            d_max=d_max,
+            d_min=d_min,
+            gate_max=gate_max,
+            scale_log2=self.softmax_scale_log2,
+            gate_threshold_log2=gate_threshold_log2,
+        )

build/torch-cuda/seqlen_info.py ADDED Viewed

	@@ -0,0 +1,273 @@

+import triton
+import triton.language as tl
+@triton.jit
+def get_seqlen_info(
+    batch_idx,
+    seqlen_static,
+    cu_seqlens,
+    seqused,
+    HAS_CU_SEQLENS: tl.constexpr,
+    HAS_SEQUSED: tl.constexpr,
+):
+    """
+    Get offset and seqlen for a given batch index.
+    :param batch_idx: Index of the batch.
+    :param seqlen_static: Static sequence length if cu_seqlens is not provided.
+    :param cu_seqlens: Cumulative sequence lengths tensor.
+    :param seqused: Actual sequence lengths tensor.
+    :param HAS_CU_SEQLENS: Boolean flag indicating if cu_seqlens is provided.
+    :param HAS_SEQUSED: Boolean flag indicating if seqused is provided.
+    :return offset: Offset for the given batch index.
+    :return seqlen: Sequence length for the given batch index.
+    """
+    if HAS_CU_SEQLENS:
+        offset = tl.load(cu_seqlens + batch_idx)
+        if HAS_SEQUSED:
+            seqlen = tl.load(seqused + batch_idx)
+        else:
+            seqlen = tl.load(cu_seqlens + batch_idx + 1) - offset
+    else:
+        offset = 0
+        if HAS_SEQUSED:
+            seqlen = tl.load(seqused + batch_idx)
+        else:
+            seqlen = seqlen_static
+    return offset, seqlen
+@triton.jit
+def get_seqlen_info_qk(
+    batch_idx,
+    seqlen_q_static,
+    seqlen_k_static,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    seqused_q,
+    seqused_k,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+    HAS_CU_SEQLENS_Q: tl.constexpr,
+    HAS_CU_SEQLENS_K: tl.constexpr,
+    HAS_SEQUSED_Q: tl.constexpr,
+    HAS_SEQUSED_K: tl.constexpr,
+):
+    """
+    Get offset, padded_offset, and seqlen for both Q and K.
+    :param batch_idx: Index of the batch.
+    :param seqlen_q_static: Static sequence length for Q if cu_seqlens_q is not provided.
+    :param seqlen_k_static: Static sequence length for K if cu_seqlens_k is not provided.
+    :param cu_seqlens_q: Cumulative sequence lengths tensor for Q.
+    :param cu_seqlens_k: Cumulative sequence lengths tensor for K.
+    :param seqused_q: Actual sequence lengths tensor for Q.
+    :param seqused_k: Actual sequence lengths tensor for K.
+    :param TILE_M: Tile size for Q.
+    :param TILE_N: Tile size for K.
+    :param HAS_CU_SEQLENS_Q: Boolean flag indicating if cu_seqlens_q is provided.
+    :param HAS_CU_SEQLENS_K: Boolean flag indicating if cu_seqlens_k is provided.
+    :param HAS_SEQUSED_Q: Boolean flag indicating if seqused_q is provided.
+    :param HAS_SEQUSED_K: Boolean flag indicating if seqused_k is provided.
+    :return offset_q: Offset for Q for the given batch index.
+    :return offset_k: Offset for K for the given batch index.
+    :return padded_offset_q: Padded offset for Q aligned to TILE_M.
+    :return padded_offset_k: Padded offset for K aligned to TILE_N.
+    :return seqlen_q: Sequence length for Q for the given batch index.
+    :return seqlen_k: Sequence length for K for the given batch index.
+    """
+    # Q offset and seqlen
+    if HAS_CU_SEQLENS_Q:
+        offset_q = tl.load(cu_seqlens_q + batch_idx)
+        padded_offset_q = (offset_q + batch_idx * TILE_M) // TILE_M * TILE_M
+        if HAS_SEQUSED_Q:
+            seqlen_q = tl.load(seqused_q + batch_idx)
+        else:
+            seqlen_q = tl.load(cu_seqlens_q + batch_idx + 1) - offset_q
+    else:
+        offset_q = 0
+        padded_offset_q = 0
+        if HAS_SEQUSED_Q:
+            seqlen_q = tl.load(seqused_q + batch_idx)
+        else:
+            seqlen_q = seqlen_q_static
+    # K offset and seqlen
+    if HAS_CU_SEQLENS_K:
+        offset_k = tl.load(cu_seqlens_k + batch_idx)
+        padded_offset_k = (offset_k + batch_idx * TILE_N) // TILE_N * TILE_N
+        if HAS_SEQUSED_K:
+            seqlen_k = tl.load(seqused_k + batch_idx)
+        else:
+            seqlen_k = tl.load(cu_seqlens_k + batch_idx + 1) - offset_k
+    else:
+        offset_k = 0
+        padded_offset_k = 0
+        if HAS_SEQUSED_K:
+            seqlen_k = tl.load(seqused_k + batch_idx)
+        else:
+            seqlen_k = seqlen_k_static
+    return offset_q, offset_k, padded_offset_q, padded_offset_k, seqlen_q, seqlen_k
+@triton.jit
+def get_softmax_threshold(
+    softmax_threshold,
+    m_block,
+    seqlen_q,
+    seqlen_k,
+    IS_CAUSAL: tl.constexpr,
+    TILE_M: tl.constexpr,
+    QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr,
+):
+    """
+    Compute the softmax threshold for a given block.
+    :param softmax_threshold: Threshold value normalized by full key length.
+    :param m_block: Current block index along the M dimension.
+    :param seqlen_q: Sequence length of the query.
+    :param seqlen_k: Sequence length of the key.
+    :param IS_CAUSAL: Boolean flag indicating if the attention is causal.
+    :param TILE_M: Tile size along the M dimension.
+    :param QHEAD_PER_KVHEAD_PACKGQA: Ratio of query heads to key/value heads for packed GQA.
+    :return softmax_threshold_log2: softmax threshold of shape [TILE_M] in log2-domain for the given block.
+    """
+    if IS_CAUSAL:
+        offs_m = m_block * TILE_M + tl.arange(0, TILE_M)
+        q_idx = offs_m
+        if QHEAD_PER_KVHEAD_PACKGQA > 1:
+            q_idx = q_idx // QHEAD_PER_KVHEAD_PACKGQA
+        causal_offset = seqlen_k - seqlen_q
+        s_thr = softmax_threshold * (q_idx + causal_offset + 1.0) / seqlen_k
+    else:
+        s_thr = tl.full((TILE_M,), softmax_threshold, dtype=tl.float32)
+    softmax_threshold_log2 = tl.log2(s_thr)
+    return softmax_threshold_log2
+@triton.jit
+def get_gate_threshold(
+    gate_threshold,
+    m_block,
+    seqlen_q,
+    seqlen_k,
+    IS_CAUSAL: tl.constexpr,
+    TILE_M: tl.constexpr,
+    QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr,
+    IS_ADAPT_GATE: tl.constexpr,
+):
+    """
+    Compute the gate threshold for a given block.
+    :param gate_threshold: Threshold value for the gate.
+    :param m_block: Current block index along the M dimension.
+    :param seqlen_q: Sequence length of the query.
+    :param seqlen_k: Sequence length of the key.
+    :param IS_CAUSAL: Boolean flag indicating if the attention is causal.
+    :param TILE_M: Tile size along the M dimension.
+    :param QHEAD_PER_KVHEAD_PACKGQA: Ratio of query heads to key/value heads for packed GQA.
+    :param IS_ADAPT_GATE: Boolean flag indicating if self-adaptive gate threshold is enabled.
+    :return gate_threshold_log2: Lower-bound scalar gate threshold in log2-domain for the given block.
+    """
+    if IS_CAUSAL and IS_ADAPT_GATE:
+        offs_m = m_block * TILE_M + tl.arange(0, TILE_M)
+        q_idx = offs_m
+        if QHEAD_PER_KVHEAD_PACKGQA > 1:
+            q_idx = q_idx // QHEAD_PER_KVHEAD_PACKGQA
+        causal_offset = seqlen_k - seqlen_q
+        g_thr = tl.min(gate_threshold * (q_idx + causal_offset + 1.0) / seqlen_k)
+    else:
+        g_thr = tl.full((), gate_threshold, dtype=tl.float32)
+    gate_threshold_log2 = tl.log2(g_thr)
+    return gate_threshold_log2
+@triton.jit
+def offset_batch_Q(
+    base_ptr,
+    batch_idx,
+    offset,
+    padded_offset,
+    stride_batch,
+    stride_seq,
+    HAS_CU_SEQLENS: tl.constexpr,
+    USE_PADDED: tl.constexpr,
+):
+    if HAS_CU_SEQLENS:
+        actual_offset = padded_offset if USE_PADDED else offset
+        return base_ptr + actual_offset * stride_seq
+    else:
+        return base_ptr + batch_idx * stride_batch
+@triton.jit
+def offset_batch_K(
+    base_ptr,
+    batch_idx,
+    offset,
+    padded_offset,
+    stride_batch,
+    stride_seq,
+    HAS_CU_SEQLENS: tl.constexpr,
+    USE_PADDED: tl.constexpr,
+):
+    if HAS_CU_SEQLENS:
+        actual_offset = padded_offset if USE_PADDED else offset
+        return base_ptr + actual_offset * stride_seq
+    else:
+        return base_ptr + batch_idx * stride_batch
+@triton.jit
+def make_ptrs(
+    base_ptrs,
+    mn_block,
+    stride_seq,
+    TILE_MN: tl.constexpr,
+    TILE_K: tl.constexpr,
+    SWAP_AB: tl.constexpr,
+):
+    offs_mn = mn_block * TILE_MN + tl.arange(0, TILE_MN)
+    if TILE_K > 1:
+        offs_k = tl.arange(0, TILE_K)
+        if SWAP_AB:
+            ptrs = base_ptrs + offs_mn[None, :] * stride_seq + offs_k[:, None]
+        else:
+            ptrs = base_ptrs + offs_mn[:, None] * stride_seq + offs_k[None, :]
+    else:
+        ptrs = base_ptrs + offs_mn * stride_seq
+    return ptrs
+@triton.jit
+def make_pack_gqa_ptrs(
+    base_ptrs,
+    m_block,
+    head_idx,
+    stride_head,
+    stride_seq,
+    TILE_M: tl.constexpr,
+    TILE_K: tl.constexpr,
+    QHEAD_PER_KVHEAD_PACKGQA: tl.constexpr,
+):
+    offs_m = m_block * TILE_M + tl.arange(0, TILE_M)
+    m_idx = offs_m // QHEAD_PER_KVHEAD_PACKGQA
+    q_head_offset = offs_m - m_idx * QHEAD_PER_KVHEAD_PACKGQA
+    q_head = head_idx * QHEAD_PER_KVHEAD_PACKGQA + q_head_offset
+    if TILE_K > 1:
+        offs_k = tl.arange(0, TILE_K)
+        ptrs = (
+            base_ptrs
+            + m_idx[:, None] * stride_seq
+            + q_head[:, None] * stride_head
+            + offs_k[None, :]
+        )
+    else:
+        ptrs = base_ptrs + m_idx * stride_seq + q_head * stride_head
+    return ptrs

build/torch-cuda/utils.py ADDED Viewed

	@@ -0,0 +1,167 @@

+import functools
+import torch
+import triton
+def get_device():
+    """
+    Get the appropriate device for computation.
+    :return device: torch.device object
+    """
+    # TODO: add NPU
+    # Works for both NVIDIA and AMD
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    # Intel XPU if available
+    elif torch.xpu.is_available():
+        return torch.device("xpu")
+    elif torch.mps.is_available():
+        return torch.device("mps")
+    else:
+        return torch.device("cpu")
+def ensure_contiguous(fn):
+    """
+    Decorator to ensure that all tensor inputs to the decorated function are contiguous.
+    :param fn: Function to be decorated
+    :return wrapper: Wrapped function
+    """
+    @functools.wraps(fn)
+    def wrapper(ctx, *args, **kwargs):
+        def maybe_to_contiguous(x):
+            return x.contiguous() if isinstance(x, torch.Tensor) else x
+        args = [maybe_to_contiguous(arg) for arg in args]
+        kwargs = {k: maybe_to_contiguous(v) for k, v in kwargs.items()}
+        return fn(ctx, *args, **kwargs)
+    return wrapper
+@functools.lru_cache(maxsize=4096)
+def window_sizes_heuristic(
+    seqlen_k: int,
+    num_heads_kv: int,
+    device: torch.device,
+    equal_bandwidth: bool = True,
+    window_dist: int = 1024,
+    window_sink: int = 64,
+) -> torch.Tensor:
+    """
+    Compute token count window sizes that partition non-diagonal causal distances into bands.
+    :param seqlen_k: Sequence length of keys.
+    :param num_heads_kv: Number of KV heads.
+    :param device: Target device.
+    :param equal_bandwidth: If True, use equal-bandwidth partitioning for balanced decode load. If False, use equal-area partitioning for balanced forward and backward load.
+    :param window_dist: Near-diagonal token count shared by all KV heads.
+    :param window_sink: Sink token count shared by all KV heads.
+    :return: int32 tensor with shape [num_heads_kv, 4], columns are [window_sink, window_left, window_right, window_dist]. window_sink is the prefix sink token count, window_left is the distant local band token count, window_right is the token count gap after the near-diagonal window before the distant band, and window_dist is the near-diagonal token count.
+    """
+    window_dist = min(max(window_dist, 0), seqlen_k)
+    window_sink = min(max(window_sink, 0), max(seqlen_k - window_dist, 0))
+    head_kv_idx = torch.arange(num_heads_kv + 1, dtype=torch.float32)
+    distance_span = max(seqlen_k - window_sink - window_dist, 0)
+    if equal_bandwidth:
+        breakpoints = (distance_span * head_kv_idx / num_heads_kv).to(torch.int32)
+    else:
+        breakpoints = (
+            distance_span * (1.0 - torch.sqrt(1.0 - head_kv_idx / num_heads_kv))
+        ).to(torch.int32)
+    window_size_left = breakpoints[1:] - breakpoints[:-1]
+    window_size_right = breakpoints[:-1]
+    window_size_dist = torch.full_like(window_size_left, window_dist)
+    window_size_sink = torch.full_like(window_size_left, window_sink)
+    return torch.stack(
+        [window_size_sink, window_size_left, window_size_right, window_size_dist], dim=1
+    ).to(device)
+def max_split_blocks_from_window_sizes(
+    window_sizes: torch.Tensor, TILE_N: int
+) -> int | None:
+    """
+    Estimate useful split-axis blocks for local attention.
+    The exact useful split count depends on the max values stored in window_sizes.
+    :param window_sizes: Int32 tensor with shape [num_heads_kv, 4], columns are [sink, left, right, dist]
+    :param TILE_N: Tile size for N dimension.
+    :return: Useful split-axis blocks, or None when CUDA Graph capture needs a static launch-grid bound.
+    """
+    if window_sizes is None or window_sizes.numel() == 0:
+        return 1
+    if (
+        window_sizes.is_cuda
+        and torch.cuda.is_available()
+        and torch.cuda.is_current_stream_capturing()
+    ):
+        return None
+    window_sizes = torch.clamp(window_sizes, min=0)
+    window_sink = torch.max(window_sizes[:, 0])
+    window_left = torch.max(window_sizes[:, 1])
+    window_dist = torch.max(window_sizes[:, 3])
+    distant_blocks = torch.div(window_left + TILE_N - 1, TILE_N, rounding_mode="floor")
+    near_blocks = (window_dist > 0).to(window_left.dtype)
+    sink_blocks = torch.div(window_sink + TILE_N - 1, TILE_N, rounding_mode="floor")
+    split_blocks = torch.maximum(
+        distant_blocks + near_blocks + sink_blocks,
+        torch.ones((), dtype=window_left.dtype, device=window_sizes.device),
+    )
+    return int(split_blocks)
+@functools.lru_cache(maxsize=4096)
+def num_splits_heuristic(
+    seqlen_q: int,
+    seqlen_k: int,
+    num_SMs: int,
+    TILE_M: int,
+    TILE_N: int,
+    max_split_blocks: int | None = None,
+) -> int:
+    """
+    Determine the number of KV splits for FlashDecoding.
+    Splits only when there are enough KV blocks to benefit from parallelism,
+    and targets full SM occupancy by over-subscribing the M-block count.
+    :param seqlen_q: Sequence length of queries.
+    :param seqlen_k: Sequence length of keys.
+    :param num_SMs: Number of streaming multiprocessors on the device.
+    :param TILE_M: Tile size for M dimension.
+    :param TILE_N: Tile size for N dimension.
+    :param max_split_blocks: Optional cap on useful split-axis blocks. Local two-band attention should pass this to avoid assigning splits to empty gap regions.
+    :return: Number of splits.
+    """
+    total_mblocks = triton.cdiv(seqlen_q, TILE_M)
+    num_n_blocks = triton.cdiv(seqlen_k, TILE_N)
+    effective_n_blocks = num_n_blocks
+    if max_split_blocks is not None:
+        effective_n_blocks = min(num_n_blocks, max(max_split_blocks, 1))
+    max_splits = 1 << (max(num_SMs, 1).bit_length() - 1)
+    if effective_n_blocks <= 4:
+        # 1 means no splitting
+        return 1
+    return max(1, min(num_SMs // max(total_mblocks, 1), max_splits, effective_n_blocks))
+def alloc_fn(size: int, alignment: int, stream):
+    """
+    TMA descriptors require a global memory allocation
+    :param size: Size of the allocation in bytes.
+    :param alignment: Alignment requirement in bytes.
+    :param stream: CUDA stream for the allocation.
+    :return: A torch.Tensor representing the allocated memory.
+    """
+    return torch.empty(size, device="cuda", dtype=torch.int8)