# mypy: allow-untyped-defs import logging from typing import Optional, Tuple import torch import torch.nn import torch.nn.functional as F from torch.backends.cuda import ( can_use_efficient_attention, can_use_flash_attention, flash_sdp_enabled, math_sdp_enabled, mem_efficient_sdp_enabled, SDPAParams, ) from torch.nn.attention import SDPBackend from .nested_tensor import NestedTensor log = logging.getLogger(__name__) def _validate_sdpa_input( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tensor] = None, dropout_p=0.0, is_causal=False, scale=None, ): if ( not isinstance(query, NestedTensor) or not isinstance(key, NestedTensor) or not isinstance(value, NestedTensor) ): raise ValueError( f"Expected query, key, and value to be nested tensors, " f"but got query.is_nested: {query.is_nested}, key.is_nested: {key.is_nested}, " f"and value.is_nested: {value.is_nested} instead." ) if query.dtype != key.dtype or query.dtype != value.dtype: raise ValueError( f"Expected query, key, and value to have the same dtype, " f"but got query.dtype: {query.dtype}, key.dtype: {key.dtype}, " f"and value.dtype: {value.dtype} instead." ) if query.device != key.device or query.device != value.device: raise ValueError( f"Expected query, key, and value to have the same device type, " f"but got query.device: {query.device}, key.device: {key.device}, " f"and value.device: {value.device} instead." ) if query.dim() < 3 or key.dim() < 3 or value.dim() < 3: raise ValueError( f"Expected query, key, and value to all be at least 3 dimensional, but got query.dim: " f"{query.dim()}, key.dim: {key.dim()} and value.dim: {value.dim()} instead." ) if query._ragged_idx != key._ragged_idx or query._ragged_idx != value._ragged_idx: raise ValueError( f"Expected query, key, and value to all be ragged on the same dimension, but got ragged " f"dims {query._ragged_idx}, {key._ragged_idx}, and {value._ragged_idx}, respectively." ) if attn_mask is not None: # TODO: Figure out whether masks are actually supported for this layout or not raise ValueError("Masks are not yet supported!") if attn_mask.dtype != torch.bool and attn_mask.dtype != query.dtype: raise ValueError( f"Expected attn_mask dtype to be bool or to match query dtype, but got attn_mask.dtype: " f"{attn_mask.dtype}, and query.dtype: {query.dtype} instead." ) def _check_batch_size_nested(params: SDPAParams, debug=False) -> bool: # This is expected to be called after check_tensor_shapes ensuring that the # size() calls won't error since the inputs are all 4 dimensional q_batch_size = params.query.size(0) k_batch_size = params.key.size(0) v_batch_size = params.value.size(0) # num_heads logic for nested input is checked in # check_for_seq_len_0_nested_tensor as there is handling there to make sure # num_heads is not ragged return q_batch_size == k_batch_size and q_batch_size == v_batch_size def _check_head_dim_size_flash_nested(params: SDPAParams, debug=False) -> bool: max_size = 256 query_size_last = params.query.size(-1) key_size_last = params.key.size(-1) value_size_last = params.value.size(-1) same_head_dim_size = ( query_size_last == key_size_last and query_size_last == value_size_last ) if not ( same_head_dim_size and (query_size_last % 8 == 0) and (query_size_last <= max_size) ): if debug: log.warning( "For NestedTensor inputs, Flash attention requires q,k,v to have the same " "last dimension and to be a multiple of 8 and less than or equal to 256. " "Got Query.size(-1): %d, Key.size(-1): %d, Value.size(-1): %d instead.", query_size_last, key_size_last, value_size_last, ) return False return True def _check_for_seq_len_0_and_consistent_head_dim_nested_helper( param: torch.Tensor, param_name: str, debug=False ) -> bool: assert isinstance(param, NestedTensor), "param should be a jagged NT" if param._ragged_idx == 1: # num_head_dims is ragged if debug: log.warning( "Fused kernels do not support ragged num_head_dims, %s has a ragged num_heads.", param_name, ) return False # This is being called inside sdp with shape [batch, heads, {seq_len}, dim] if param._get_min_seqlen() == 0: if debug: log.warning( "Fused kernels do not support seq_len == 0, %s has a seq len of 0.", param_name, ) return False return True def _try_broadcast_param_size(q_size, k_size, v_size, param_name, debug=False) -> bool: max_size = max(q_size, k_size, v_size) if ( (q_size != max_size and q_size != 1) or (k_size != max_size and k_size != 1) or (v_size != max_size and v_size != 1) ): if debug: log.warning( "Both fused kernels require query, key and value to have broadcastable %s, " "got Query %s %d, Key %s %d, Value %s %d instead.", param_name, param_name, q_size, param_name, k_size, param_name, v_size, ) return False return True def _check_for_seq_len_0_nested(params: SDPAParams, debug=False) -> bool: # When this function is called we are assured that the nt is dim==4 q_is_safe = ( _check_for_seq_len_0_and_consistent_head_dim_nested_helper( params.query, "query", debug ) if params.query.is_nested else True ) # short circuit if any is unsafe if not q_is_safe: return False k_is_safe = ( _check_for_seq_len_0_and_consistent_head_dim_nested_helper( params.key, "key", debug ) if params.key.is_nested else True ) # short circuit if any is unsafe if not k_is_safe: return False v_is_safe = ( _check_for_seq_len_0_and_consistent_head_dim_nested_helper( params.value, "value", debug ) if params.value.is_nested else True ) # short circuit if any is unsafe if not v_is_safe: return False # We now know none of the inputs have ragged num_heads, so we can safely # access .size(1) q_num_heads = params.query.size(1) k_num_heads = params.key.size(1) v_num_heads = params.value.size(1) same_num_heads = q_num_heads == k_num_heads and q_num_heads == v_num_heads if not same_num_heads: if ( params.query.requires_grad or params.key.requires_grad or params.value.requires_grad ): if debug: log.warning( "Both fused kernels do not support training with broadcasted NT inputs." ) return False return _try_broadcast_param_size( q_num_heads, k_num_heads, v_num_heads, "num heads", debug ) return True def _can_use_flash_sdpa_jagged(params: SDPAParams, debug=False) -> bool: constraints = ( _check_batch_size_nested, _check_head_dim_size_flash_nested, _check_for_seq_len_0_nested, ) for constraint in constraints: if not constraint(params, debug): return False return True def _can_use_efficient_sdpa_jagged(params: SDPAParams, debug=False) -> bool: constraints = ( _check_batch_size_nested, _check_for_seq_len_0_nested, ) for constraint in constraints: if not constraint(params, debug): return False return True def _can_use_math_sdpa_jagged(params: SDPAParams, debug=False) -> bool: if ( not params.query.transpose(1, 2).is_contiguous() or not params.key.transpose(1, 2).is_contiguous() or not params.value.transpose(1, 2).is_contiguous() ): if debug: log.warning( "If inputs are nested tensors they must be contiguous after transposing." ) return False if params.is_causal: if debug: log.warning( "Nested tensors for query / key are not supported when is_causal=True." ) return False return True def _select_sdp_backend(query, key, value, attn_mask, dropout, is_causal, enable_gqa): if ( not flash_sdp_enabled() and not mem_efficient_sdp_enabled() and not math_sdp_enabled() ): return SDPBackend.ERROR ordering = ( SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH, ) params = SDPAParams(query, key, value, attn_mask, dropout, is_causal, enable_gqa) for backend in ordering: if backend == SDPBackend.FLASH_ATTENTION: if can_use_flash_attention(params) and _can_use_flash_sdpa_jagged(params): return SDPBackend.FLASH_ATTENTION if backend == SDPBackend.EFFICIENT_ATTENTION: if can_use_efficient_attention(params) and _can_use_efficient_sdpa_jagged( params ): return SDPBackend.EFFICIENT_ATTENTION if backend == SDPBackend.MATH: if math_sdp_enabled() and _can_use_math_sdpa_jagged(params): return SDPBackend.MATH log.warning("Memory efficient kernel not used because:") can_use_efficient_attention(params, debug=True) _can_use_efficient_sdpa_jagged(params, debug=True) log.warning("Flash attention kernel not used because:") can_use_flash_attention(params, debug=True) _can_use_flash_sdpa_jagged(params, debug=True) log.warning("Math attention kernel not used because:") _can_use_math_sdpa_jagged(params, debug=True) return SDPBackend.ERROR def _cumulative_and_max_seq_len_nnz(qkv: torch.Tensor) -> Tuple[torch.Tensor, int, int]: # This function is used to calculate two pieces of metadata that are needed # for use with flash-attention and efficient_attention kernels. They are the # cumulative sequence_length over a batch of sequences and the maximum # sequence length. # It returns a tuple of cumulative sequence lengths and the maximum sequence # length, and the last element in the cumulative_sequence_lengths if not isinstance(qkv, NestedTensor): raise ValueError("QKV must be nested for flash cumulative_seq_len calculation.") if qkv.lengths() is None: # TODO: Explore performance impact of copying cumulative_seqlen = qkv.offsets().to(dtype=torch.int32, device=qkv.device) max_seqlen = qkv._get_max_seqlen() n_elem = qkv.values().shape[0] else: # TODO: Explore performance impact of copying cumulative_seqlen = ( qkv.lengths().cumsum(0).to(dtype=torch.int32, device=qkv.device) ) batch_size = qkv.size(0) max_seqlen = qkv._get_max_seqlen() # TODO: Explore performance impact when compiling n_elem = int(cumulative_seqlen[-1].item()) return cumulative_seqlen, max_seqlen, n_elem def _is_safe_to_get_storage_as_tensor(tensor: torch.Tensor): # This function checks if a nested tensor is valid for # use with the flash-attention and efficient_attention kernels without # needing to call contiguous on the nested tensor input. # It checks that the storage offsets' adjacent_differences are a constant # mutiple of the previous tensor in the nested tensor and that the strides # are monitonically decreasing. This check is done after calling transpose on # the nested tensor resulting in a Nt of shape [bsz, {seq_len}, num_heads, dim] # Returns a boolean indicating if contiguous needs to be called for input assert isinstance(tensor, NestedTensor) offsets = tensor.offsets() strides = tensor._strides n_tensors = offsets.size(0) - 1 if n_tensors <= 1: return True # Check initially that the tensor strides are in strictly descending order prev_stride = strides[1] for stride in strides[2:]: if prev_stride <= stride: # This would mean that the last stride is greater than the seq_len # stride return False prev_stride = stride # Congrats you made it! return True def _view_as_dense( tensor: torch.Tensor, Nnz: int, num_heads: int, head_dim: int ) -> torch.Tensor: if tensor.is_nested: return tensor.values() return tensor.view(Nnz, num_heads, head_dim) # TODO: Next iteration should add test cases and check it works # def _sdpa_nested_preprocessing_with_broadcast(query, key, value): # # Query (Batch x Num_heads x {Q_seq_len} x Dim_per_head) # # Key (Batch x Num_heads x {KV_seq_len} x Dim_per_head) # # Value (Batch x Num_heads x {KV_seq_len} x Dim_per_head) # q_batch_size = query.size(0) # k_batch_size = key.size(0) # v_batch_size = value.size(0) # output_batch_size = max(q_batch_size, k_batch_size, v_batch_size) # q_num_heads = query.size(1) # k_num_heads = key.size(1) # v_num_heads = value.size(1) # output_num_heads = max(q_num_heads, k_num_heads, v_num_heads) # head_dim_qk = query.size(3) # head_dim_v = value.size(3) # q_t = query.transpose(1, 2) # k_t = key.transpose(1, 2) # v_t = value.transpose(1, 2) # # Checks in sdp_utils ensure that if {*}_batch_size/{*}_num_heads != # # output_batch_size/num_heads then they are 1 # q_batch_size_needs_broadcast = q_batch_size != output_batch_size # k_batch_size_needs_broadcast = k_batch_size != output_batch_size # v_batch_size_needs_broadcast = v_batch_size != output_batch_size # # If {*}_batch_size_needs_broadcast, then # # (1) max_seqlen_batch_{*} is given by {*}_t.size(1) # # this is because needs_broadcast indicates that the batch_size is 1 # # and hence there is only 1 value for seq_len # # (2) The cum_seq_lens are given by [0, {*}_t.size(1), 2 * {*}_t.size(1), # # ..., outut_batch_size * {*}_t.size(1)] # # (3) Nnz_{*} is given by output_batch_size * {*}_t.size(1) # if q_batch_size_needs_broadcast or not q_t.is_nested: # max_seqlen_batch_q = q_t.size(1) # cumulative_sequence_length_q = torch.arange( # 0, # (output_batch_size + 1) * max_seqlen_batch_q, # max_seqlen_batch_q, # device=q_t.device, # dtype=torch.int32, # ) # Nnz_q = output_batch_size * max_seqlen_batch_q # else: # ( # cumulative_sequence_length_q, # max_seqlen_batch_q, # Nnz_q, # ) = _cumulative_and_max_seq_len_nnz(q_t) # if k_batch_size_needs_broadcast and v_batch_size_needs_broadcast: # assert k_t.size(1) == v_t.size(1) # max_seqlen_batch_kv = k_t.size(1) # cumulative_sequence_length_kv = torch.arange( # 0, # (output_batch_size + 1) * max_seqlen_batch_kv, # max_seqlen_batch_kv, # device=k_t.device, # dtype=torch.int32, # ) # Nnz_kv = output_batch_size * max_seqlen_batch_kv # else: # cumulative_sequence_length_kv, max_seqlen_batch_kv, Nnz_kv = ( # _cumulative_and_max_seq_len_nnz(v_t) # if k_batch_size_needs_broadcast # else _cumulative_and_max_seq_len_nnz(k_t) # ) # q_num_heads_needs_broadcast = q_num_heads != output_num_heads # k_num_heads_needs_broadcast = k_num_heads != output_num_heads # v_num_heads_needs_broadcast = v_num_heads != output_num_heads # if not q_t.is_nested: # query_buffer_reshaped = q_t.expand( # output_batch_size, q_t.size(1), output_num_heads, head_dim_qk # ) # query_buffer_reshaped = query_buffer_reshaped.reshape( # Nnz_q, output_num_heads, head_dim_qk # ) # else: # if not q_t.is_contiguous() and not _is_safe_to_get_storage_as_tensor(q_t): # q_t = q_t.contiguous() # # If we are broadcasting then Nnz_q will be the output_batch_size since # # seq_len is 1 # effective_batch_size_q = ( # output_batch_size if q_batch_size_needs_broadcast else Nnz_q # ) # query_buffer_reshaped = _view_as_dense( # q_t, effective_batch_size_q, output_num_heads, head_dim_qk # ) # # If the physical layout of the NestedTensor's storage # # is not: batch, {seq_len}, num_heads, head_dim then we need # # to call contiguous # if not k_t.is_contiguous() and not _is_safe_to_get_storage_as_tensor(k_t): # k_t = k_t.contiguous() # if not v_t.is_contiguous() and not _is_safe_to_get_storage_as_tensor(v_t): # v_t = v_t.contiguous() # effective_batch_size_k = ( # output_batch_size if k_batch_size_needs_broadcast else Nnz_kv # ) # key_buffer_reshaped = _view_as_dense( # k_t, effective_batch_size_k, output_num_heads, head_dim_qk # ) # effective_batch_size_v = ( # output_batch_size if v_batch_size_needs_broadcast else Nnz_kv # ) # value_buffer_reshaped = _view_as_dense( # v_t, effective_batch_size_v, output_num_heads, head_dim_v # ) # if not q_batch_size_needs_broadcast: # output_shape = q_t._size # if head_dim_v != head_dim_qk: # output_shape[-1] = head_dim_v # if q_num_heads_needs_broadcast: # output_shape[1] = output_num_heads # else: # output_shape = torch.empty(3, dtype=torch.int64, device=torch.device("cpu")) # output_shape[0] = q_t.size(1) # output_shape[1] = output_num_heads # output_shape[2] = head_dim_v # return ( # query_buffer_reshaped, # key_buffer_reshaped, # value_buffer_reshaped, # cumulative_sequence_length_q, # cumulative_sequence_length_kv, # max_seqlen_batch_q, # max_seqlen_batch_kv, # output_shape, # ) def _sdpa_nested_preprocessing(query, key, value): # Query (Batch x Num_heads x {Q_seq_len} x Dim_per_head) # Key (Batch x Num_heads x {KV_seq_len} x Dim_per_head) # Value (Batch x Num_heads x {KV_seq_len} x Dim_per_head) q_batch_size = query.size(0) k_batch_size = key.size(0) v_batch_size = value.size(0) q_num_heads = query.size(1) k_num_heads = key.size(1) v_num_heads = value.size(1) if not (q_batch_size == k_batch_size and q_batch_size == v_batch_size) or not ( q_num_heads == k_num_heads and k_num_heads == v_num_heads ): raise RuntimeError( "This path is currently not implemented for jagged layout NT." ) # return _sdpa_nested_preprocessing_with_broadcast(query, key, value) num_heads = query.size(1) head_dim_qk = query.size(3) head_dim_v = value.size(3) q_t = query.transpose(1, 2) k_t = key.transpose(1, 2) v_t = value.transpose(1, 2) ( cumulative_sequence_length_q, max_seqlen_batch_q, Nnz_q, ) = _cumulative_and_max_seq_len_nnz(q_t) ( cumulative_sequence_length_kv, max_seqlen_batch_kv, Nnz_kv, ) = _cumulative_and_max_seq_len_nnz(k_t) # [TODO] K and V have to have the same Nnz, should probably torch_check # assume in order to not iterate over v # If the physical layout of the NestedTensor's storage # is not: batch, {seq_len}, num_heads, head_dim then we need # to call contiguous if not q_t.is_contiguous() and not _is_safe_to_get_storage_as_tensor(q_t): q_t = q_t.contiguous() if not k_t.is_contiguous() and not _is_safe_to_get_storage_as_tensor(k_t): k_t = k_t.contiguous() if not v_t.is_contiguous() and not _is_safe_to_get_storage_as_tensor(v_t): v_t = v_t.contiguous() query_buffer_reshaped = _view_as_dense(q_t, Nnz_q, num_heads, head_dim_qk) key_buffer_reshaped = _view_as_dense(k_t, Nnz_kv, num_heads, head_dim_qk) value_buffer_reshaped = _view_as_dense(v_t, Nnz_kv, num_heads, head_dim_v) output_nt_info = { "offsets": q_t.offsets(), "_max_seqlen": q_t._get_max_seqlen(), "_min_seqlen": q_t._get_min_seqlen(), } return ( query_buffer_reshaped, key_buffer_reshaped, value_buffer_reshaped, cumulative_sequence_length_q, cumulative_sequence_length_kv, max_seqlen_batch_q, max_seqlen_batch_kv, output_nt_info, ) def _pad_last_dim( tensor: torch.Tensor, alignment_size: int, slice: bool ) -> torch.Tensor: # FlashAttentionV2 requires that head dimension be a multiple of 8 # This was previously done within the kernel, however # This causes the kernel to maybe alias query, key, value # So instead we pad the head_dimensions to be a multiple of 8 # in the composite region last_dim_size = tensor.size(-1) if last_dim_size % alignment_size == 0: return tensor pad_count = alignment_size - (last_dim_size % alignment_size) tensor = torch.nn.functional.pad(tensor, [0, pad_count]) if slice: return tensor[..., 0:last_dim_size] return tensor # TODO: coalesce with torch/nn/utils/attention.py def _calculate_scale(query, scale): # TODO: Investigate why math.sqrt() isn't properly handled by Dynamo? softmax_scale = scale if scale is not None else torch.sym_sqrt(1.0 / query.size(-1)) return softmax_scale def _post_process_flash_output(out: torch.Tensor, og_size): if not out.is_nested and out.size(-1) != og_size: out = out[..., 0:og_size] return out def _is_computing_meta_flops(x): # Note: there's a use case of using meta tensors & the dispatch-based flop counter. # We can use this function to check for this scenario in order to handle it specially. if not torch.jit.is_scripting() and x.device.type == "meta": torch_dispatch_mode_stack = ( torch.utils._python_dispatch._get_current_dispatch_mode_stack() ) return any( type(x) == torch.utils.flop_counter.FlopCounterMode for x in torch_dispatch_mode_stack ) return False def _autocast( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tensor], ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: """ [Autocasting SDPA for NJT] Normal autocasting doesn't work for NJT+SDPA right now: * NJT intercepts the __torch_function__ call for scaled_dot_product_attention, which happens before we get to any aten ops or dispatcher logic; then the torch_function logic calls into efficient attention or flash attention. So, autocasting on the scaled_dot_product_attention op won't work because we never see that aten op. * If we put autocasting on `_flash_attention_forward`, then we'll get autocasting to run, but the kernel selection logic in torch_function handling (ie. jagged_scaled_dot_product_attention) won't work correctly: the kernel selection logic will run before autocasting, and choose a kernel based on the un-autocasted dtypes; but then autocasting will run and the actual attention computation will happen in a different dtype. An alternative is to just change the backend selection logic for SDPA+NJT to be autocast-aware and rely on autocasting to do the actual conversions for flash attention / efficient attention. However, by manually doing the actual autocast before the backend selection, we ensure that the autocast handling for backend selection doesn't diverge from the autocast handling for the actual dtype conversions. """ device_type = query.device.type # meta device is not supported by autocast, so break early for it if _is_computing_meta_flops(query) or not torch.is_autocast_enabled(device_type): return query, key, value, attn_mask def cvt(x): if x is None: return x target_dtype = torch.get_autocast_dtype(device_type) if ( (not x.dtype.is_floating_point) or x.dtype == target_dtype or x.dtype == torch.float64 ): return x return x.to(target_dtype) return cvt(query), cvt(key), cvt(value), cvt(attn_mask) def jagged_scaled_dot_product_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tensor] = None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False, ): query, key, value, attn_mask = _autocast(query, key, value, attn_mask) _validate_sdpa_input(query, key, value, attn_mask, dropout_p, is_causal, scale) # for mypy, ugh assert ( isinstance(query, NestedTensor) and isinstance(key, NestedTensor) and isinstance(value, NestedTensor) ) from torch.nested._internal.nested_tensor import nested_view_from_values_offsets # Special path for non-ragged sequence length (e.g. for SAM where we have a ragged # second batch dim instead). For this case, we can just send the dense buffers through # vanilla SDPA. if query.dim() > 3 and key.dim() > 3 and value.dim() > 3 and query._ragged_idx == 1: output = F.scaled_dot_product_attention( query.values(), key.values(), value.values(), attn_mask=( attn_mask.values() if isinstance(attn_mask, NestedTensor) else attn_mask ), dropout_p=dropout_p, is_causal=is_causal, scale=scale, ) return nested_view_from_values_offsets(output, query.offsets()) compute_logsumexp = query.requires_grad or key.requires_grad or value.requires_grad backend_choice = _select_sdp_backend( query, key, value, attn_mask, dropout_p, is_causal, enable_gqa ) if _is_computing_meta_flops(query): # Backend choice will probably not be correct if we have a meta device, # because backend choice is device-aware. In this case, we mostly just # want to avoid using math backend (which does a .item() call). # Arbitrarily choose flash attention. backend_choice = SDPBackend.FLASH_ATTENTION if backend_choice == SDPBackend.FLASH_ATTENTION: og_size = query.size(-1) query_padded = _pad_last_dim(query, 8, False) key_padded = _pad_last_dim(key, 8, False) value_padded = _pad_last_dim(value, 8, False) # We need to calculate the scale based off the OG head dim size og_scale = _calculate_scale(query, scale) ( query_buffer_reshaped, key_buffer_reshaped, value_buffer_reshaped, cumulative_sequence_length_q, cumulative_sequence_length_kv, max_seqlen_batch_q, max_seqlen_batch_kv, output_nt_info, ) = _sdpa_nested_preprocessing(query_padded, key_padded, value_padded) ( attention, logsumexp, philox_seed, philox_offset, debug_attn_mask, ) = torch.ops.aten._flash_attention_forward( query_buffer_reshaped, key_buffer_reshaped, value_buffer_reshaped, cumulative_sequence_length_q, cumulative_sequence_length_kv, max_seqlen_batch_q, max_seqlen_batch_kv, dropout_p, is_causal, False, scale=og_scale, ) # Reshape output to convert nnz to batch_size and seq_len attention = nested_view_from_values_offsets( attention, # output from flash_attn is [total_q, num_heads, head_size_og] output_nt_info["offsets"], min_seqlen=output_nt_info["_min_seqlen"], max_seqlen=output_nt_info["_max_seqlen"], ).transpose(1, 2) return _post_process_flash_output(attention, og_size) elif backend_choice == SDPBackend.EFFICIENT_ATTENTION: ( query_reshaped, key_reshaped, value_reshaped, cumulative_sequence_length_q, cumulative_sequence_length_kv, max_seqlen_batch_q, max_seqlen_batch_kv, output_nt_info, ) = _sdpa_nested_preprocessing(query, key, value) ( attention, log_sumexp, seed, offset, max_seqlen_q, max_seqlen_batch_kv, ) = torch.ops.aten._efficient_attention_forward( query_reshaped.unsqueeze(0), key_reshaped.unsqueeze(0), value_reshaped.unsqueeze(0), None, cumulative_sequence_length_q, cumulative_sequence_length_kv, max_seqlen_batch_q, max_seqlen_batch_kv, dropout_p, int(is_causal), compute_logsumexp, scale=scale, ) # Reshape output to convert nnz to batch_size and seq_len return nested_view_from_values_offsets( attention.squeeze(0), output_nt_info["offsets"], min_seqlen=output_nt_info["_min_seqlen"], max_seqlen=output_nt_info["_max_seqlen"], ).transpose(1, 2) elif backend_choice == SDPBackend.MATH: # save the offsets and shape of the inputs, so we can reshape the final output # query @ key = attn: [B, D1, j0, D'] @ [B, D1, D' j1] = [B, D1, j0, j1] # attn @ value = out: [B, D1, j0, j1] @ [B, D1, j1, D2] = [B, D1, j0, D2] offsets = query.offsets() d1 = query._size[1] d2 = value._size[-1] min_seqlen_tensor = query._metadata_cache.get( "min_seqlen", None ) # type: ignore[attr-defined] max_seqlen_tensor = query._metadata_cache.get( "max_seqlen", None ) # type: ignore[attr-defined] # convert jagged layout Nested Tensor to strided layout Nested Tensor # which support the math implementation of SDPA def get_strided_layout_nested_tensor(jagged_layout_nt): lengths = jagged_layout_nt._offsets[1:] - jagged_layout_nt._offsets[:-1] transpose = torch.transpose(jagged_layout_nt, 1, 2) tensor_list = transpose.values().split(list(lengths), dim=0) strided_nt = torch.nested.as_nested_tensor(list(tensor_list)) strided_nt = strided_nt.transpose(1, 2).contiguous() return strided_nt query = get_strided_layout_nested_tensor(query) key = get_strided_layout_nested_tensor(key) value = get_strided_layout_nested_tensor(value) attn_out = torch._scaled_dot_product_attention_math( query, key, value, attn_mask, dropout_p, is_causal, scale=scale )[0] from torch.nested._internal.nested_tensor import _load_val_from_tensor # convert strided layout Nested Tensor back to jagged layout Nested Tensor attn_out = attn_out.transpose(1, 2).contiguous().values() attn_out = attn_out.view(-1, d1, d2) attn_out = nested_view_from_values_offsets( attn_out, offsets, min_seqlen=( None if min_seqlen_tensor is None else _load_val_from_tensor(min_seqlen_tensor) ), max_seqlen=( None if max_seqlen_tensor is None else _load_val_from_tensor(max_seqlen_tensor) ), ).transpose(1, 2) return attn_out else: raise RuntimeError( "No viable backend for scaled_dot_product_attention was found." )