Spaces:
Running on L40S
Running on L40S
File size: 7,532 Bytes
9f818c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
"""
Imaginaire4 Attention Subpackage:
Unified implementation for all Attention implementations.
Flash Attention v2 (flash2) Backend: intermediate APIs
Only safe to import when FLASH2_SUPPORTED is True.
"""
from flash_attn.flash_attn_interface import flash_attn_func, flash_attn_varlen_func
from torch import Tensor
from cosmos_framework.model.attention.checks import assert_universal_tensor_checks
from cosmos_framework.model.attention.flash2.checks import flash2_attention_check
from cosmos_framework.model.attention.masks import CausalType
from cosmos_framework.model.attention.utils.environment import is_torch_compiling
def flash2_attention(
query: Tensor,
key: Tensor,
value: Tensor,
is_causal: bool = False,
causal_type: CausalType | None = None,
scale: float | None = None,
cumulative_seqlen_Q: Tensor | None = None,
cumulative_seqlen_KV: Tensor | None = None,
max_seqlen_Q: int | None = None,
max_seqlen_KV: int | None = None,
return_lse: bool = False,
backend_kwargs: dict | None = None,
deterministic: bool = False,
) -> Tensor | tuple[Tensor, Tensor]:
"""
Runs Flash Attention v2 on given operands (Q, K, V) with the heads-last contiguous layout
(`[batch, seqlen, heads, head_dim]`).
Parameters:
query (Tensor): 4-D query tensor, with the heads-last contiguous layout
(`[batch, seqlen, heads, head_dim]`)
key (Tensor): 4-D key tensor, with the heads-last contiguous layout
(`[batch, seqlen_kv, heads_kv, head_dim]`)
value (Tensor): 4-D value tensor, with heads-last contiguous layout
(`[batch, seqlen_kv, heads_kv, head_dim_v]`)
is_causal (bool): whether or not causal masking is enabled. Default is False.
causal_type (CausalType): causal masking mode. Choices: `CausalType.TopLeft`,
`CausalType.BottomRight`. Required when `is_causal = True`.
scale (float | None): Dot product scale (attention scale). Defaults to head_dim ** -0.5.
cumulative_seqlen_Q (Tensor | None): (varlen) Optional 1-D tensor with size `batch + 1`
indicating the cumulative sum of number of query tokens in each batch, with an
additional 0 element in the beginning. Must be passed together with
`cumulative_seqlen_KV` and `max_seqlen_{Q,KV}`.
cumulative_seqlen_KV (Tensor | None): (varlen) Optional 1-D tensor with size `batch + 1`
indicating the cumulative sum of number of key/value tokens in each batch, with an
additional 0 element in the beginning. Must be passed together with
`cumulative_seqlen_Q` and `max_seqlen_{Q,KV}`.
max_seqlen_Q (int | None): (varlen) Optional integer indicating the maximum query
sequence length in all batches. Must be passed together with `cumulative_seqlen_{Q,KV}`
and `max_seqlen_KV`.
max_seqlen_KV (int | None): (varlen) Optional integer indicating the maximum key/value
sequence length in all batches. Must be passed together with `cumulative_seqlen_{Q,KV}`
and `max_seqlen_Q`.
Other Parameters:
return_lse (bool): Whether to return the logsumexp values. Default is False.
backend_kwargs (dict | None): Key-value pair for passing arguments specific to Flash's
attention operator, if any.
deterministic (bool): Deterministic backward pass required.
Returns:
output (Tensor): 4-D output tensor, with the heads-last contiguous layout
(`[batch, seqlen, heads, head_dim_v]`).
logsumexp (Tensor): logsumexp tensor, with the heads-last contiguous layout
(`[batch, seqlen, heads, 1]`). Only returned when return_lse is True.
NOTE: this tensor is not contiguous in this backend (Flash2) and it should not be made
contiguous unless we can guarantee its results aren't merged via `merge_attentions`.
"""
is_varlen = cumulative_seqlen_Q is not None
assert_universal_tensor_checks(query, key, value)
backend_kwargs = backend_kwargs.copy() if backend_kwargs is not None else {}
# Determinism in backend_kwargs supersedes primary flag, if set to True
if "deterministic" in backend_kwargs:
deterministic = deterministic or backend_kwargs["deterministic"]
del backend_kwargs["deterministic"]
assert flash2_attention_check(
query_shape=query.shape,
key_shape=key.shape,
value_shape=value.shape,
dtype=query.dtype,
device=query.device,
requires_grad=query.requires_grad or key.requires_grad or value.requires_grad,
is_causal=is_causal,
causal_type=causal_type,
is_varlen=is_varlen,
deterministic=deterministic,
raise_error=True,
)
# This check introduces recompiles
if not is_torch_compiling():
if is_varlen and max_seqlen_Q == max_seqlen_KV == 0:
raise NotImplementedError(
"You're trying to use varlen attention with the flash2 backend and "
"an empty batch, which is not yet supported by flash2."
)
scale = scale if scale is not None else query.shape[-1] ** -0.5
if is_varlen:
assert query.shape[0] == key.shape[0] == value.shape[0] == 1
q = query.squeeze(0) # [total_tokens,H,D]
k = key.squeeze(0) # [total_tokens,Hkv,D]
v = value.squeeze(0) # [total_tokens,Hkv,Dv]
assert q.dim() == k.dim() == v.dim() == 3
out, lse_, _ = flash_attn_varlen_func(
q=query.squeeze(0),
k=key.squeeze(0),
v=value.squeeze(0),
cu_seqlens_q=cumulative_seqlen_Q,
cu_seqlens_k=cumulative_seqlen_KV,
max_seqlen_q=max_seqlen_Q,
max_seqlen_k=max_seqlen_KV,
softmax_scale=scale,
causal=is_causal,
return_attn_probs=True,
deterministic=deterministic,
**backend_kwargs,
# window_size=(-1, -1),
# dropout_p=0.0,
# softcap=0.0, # 0.0 means deactivated
# alibi_slopes=None,
# block_table=None,
)
assert out.dim() == 3 # [total_tokens,H,Dv]
assert lse_.dim() == 2 # [H,total_tokens]
output = out.unsqueeze(0) # [1,total_tokens,H,Dv]
lse = lse_.unsqueeze(0) # [1,H,total_tokens]
else:
output, lse, _ = flash_attn_func( # output: [B,N,H,Dv], lse: [B,H,N]
q=query,
k=key,
v=value,
softmax_scale=scale,
causal=is_causal,
return_attn_probs=True,
deterministic=deterministic,
**backend_kwargs,
# window_size=(-1, -1),
# dropout_p=0.0,
# softcap=0.0, # 0.0 means deactivated
# alibi_slopes=None,
)
assert isinstance(output, Tensor)
assert isinstance(lse, Tensor)
assert output.dim() == 4 # [B,N,H,Dv] or [1,total_tokens,H,Dv]
assert lse.dim() == 3 # [B,H,N] or [1,H,total_tokens]
# incorrect. All output and lse tensors passed into `merge_attentions` must have the same data
# pointer as their corresponding attention autograd ops!
lse = lse.permute(0, 2, 1) # [B,N,H] or [1,total_tokens,H]
if return_lse:
return output, lse
return output
|