lyssquant/maple-preview

1import inspect2import os3from typing import Optional, Tuple, TypedDict4 5import torch6import torch.nn.functional as F7 8 9try:10    from flash_attn_interface import flash_attn_func, flash_attn_varlen_func11except:12    from flash_attn import flash_attn_func, flash_attn_varlen_func13from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input14 15 16# Detect supported kwargs in FA317_sig = inspect.signature(flash_attn_func)18_flash_supports_window_size = "window_size" in _sig.parameters19_flash_accepts_deterministic = "deterministic" in _sig.parameters20_flash_accepts_softcap = "softcap" in _sig.parameters21 22 23def _get_unpad_data(attention_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]:24    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)25    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()26    max_seqlen_in_batch = seqlens_in_batch.max().item()27    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))28    return indices, cu_seqlens, max_seqlen_in_batch29 30 31def _upad_input(32    query_layer: torch.Tensor,33    key_layer: torch.Tensor,34    value_layer: torch.Tensor,35    attention_mask: torch.Tensor,36    query_length: int,37):38    indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)39    batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape40 41    key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k)42    value_layer = index_first_axis(43        value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k44    )45 46    if query_length == kv_seq_len:47        query_layer = index_first_axis(query_layer.reshape(batch_size * kv_seq_len, -1, head_dim), indices_k)48        cu_seqlens_q = cu_seqlens_k49        max_seqlen_in_batch_q = max_seqlen_in_batch_k50        indices_q = indices_k51    elif query_length == 1:52        max_seqlen_in_batch_q = 153        cu_seqlens_q = torch.arange(batch_size + 1, dtype=torch.int32, device=query_layer.device)54        indices_q = cu_seqlens_q[:-1]55        query_layer = query_layer.squeeze(1)56    else:57        attention_mask = attention_mask[:, -query_length:]58        query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q, *_ = unpad_input(query_layer, attention_mask)59 60    return (61        query_layer,62        key_layer,63        value_layer,64        indices_q,65        (cu_seqlens_q, cu_seqlens_k),66        (max_seqlen_in_batch_q, max_seqlen_in_batch_k),67    )68 69 70def prepare_fa3_from_position_ids(query, key, value, position_ids):71    query = query.view(-1, query.size(-2), query.size(-1))72    key = key.contiguous().view(-1, key.size(-2), key.size(-1))73    value = value.contiguous().view(-1, value.size(-2), value.size(-1))74    position_ids = position_ids.flatten()75    indices_q = torch.arange(position_ids.size(0), device=position_ids.device, dtype=torch.int32)76 77    cu_seq_lens = torch.cat(78        (79            indices_q[position_ids == 0],80            torch.tensor(position_ids.size(), device=position_ids.device, dtype=torch.int32),81        )82    )83 84    max_length = position_ids.max() + 185    return query, key, value, indices_q, (cu_seq_lens, cu_seq_lens), (max_length, max_length)86 87 88def fa_peft_integration_check(89    query: torch.Tensor,90    key: torch.Tensor,91    value: torch.Tensor,92    target_dtype: Optional[torch.dtype] = None,93):94    if target_dtype is None:95        return query, key, value96    if query.dtype == torch.float32:97        query = query.to(target_dtype)98        key = key.to(target_dtype)99        value = value.to(target_dtype)100    return query, key, value101 102 103deterministic_g = os.environ.get("FLASH_ATTENTION_DETERMINISTIC", "0") == "1"104 105 106def _flash_attention_forward(107    query_states: torch.Tensor,108    key_states: torch.Tensor,109    value_states: torch.Tensor,110    attention_mask: Optional[torch.Tensor],111    query_length: int,112    is_causal: bool,113    dropout: float = 0.0,114    position_ids: Optional[torch.Tensor] = None,115    softmax_scale: Optional[float] = None,116    sliding_window: Optional[int] = None,117    use_top_left_mask: bool = False,118    softcap: Optional[float] = None,119    deterministic: Optional[bool] = None,120    cu_seq_lens_q: Optional[torch.LongTensor] = None,121    cu_seq_lens_k: Optional[torch.LongTensor] = None,122    max_length_q: Optional[int] = None,123    max_length_k: Optional[int] = None,124    target_dtype: Optional[torch.dtype] = None,125    **kwargs,126):127    causal = is_causal if not use_top_left_mask else (is_causal and query_length != 1)128 129    flash_kwargs = {}130    if _flash_supports_window_size and sliding_window is not None and key_states.shape[1] > sliding_window:131        flash_kwargs["window_size"] = (sliding_window, 0)132    if _flash_accepts_deterministic:133        if deterministic is None:134            deterministic = deterministic_g135        flash_kwargs["deterministic"] = deterministic136 137    if attention_mask is not None:138        batch_size = query_states.shape[0]139        q_unpad, k_unpad, v_unpad, indices_q, (cu_seqlens_q, cu_seqlens_k), (max_q, max_k) = _upad_input(140            query_states, key_states, value_states, attention_mask, query_length141        )142        attn_output_unpad = flash_attn_varlen_func(143            q_unpad,144            k_unpad,145            v_unpad,146            cu_seqlens_q=cu_seqlens_q,147            cu_seqlens_k=cu_seqlens_k,148            max_seqlen_q=max_q,149            max_seqlen_k=max_k,150            # dropout_p=dropout,151            softmax_scale=softmax_scale,152            causal=causal,153            **flash_kwargs,154        )155        attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)156 157    elif position_ids is not None and (158        max_length_q is not None159        # This fucks up compile160        # or (query_length != 1 and not (torch.diff(position_ids, dim=-1) >= 0).all())161    ):162        batch_size = query_states.size(0)163        if cu_seq_lens_q is None or cu_seq_lens_k is None:164            q_unpad, k_unpad, v_unpad, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = (165                prepare_fa3_from_position_ids(query_states, key_states, value_states, position_ids)166            )167        else:168            q_unpad = query_states.reshape(-1, query_states.size(-2), query_states.size(-1))169            k_unpad = key_states.reshape(-1, key_states.size(-2), key_states.size(-1))170            v_unpad = value_states.reshape(-1, value_states.size(-2), value_states.size(-1))171 172        attn_output = flash_attn_varlen_func(173            q_unpad,174            k_unpad,175            v_unpad,176            cu_seqlens_q=cu_seq_lens_q,177            cu_seqlens_k=cu_seq_lens_k,178            max_seqlen_q=max_length_q,179            max_seqlen_k=max_length_k,180            # dropout_p=dropout,181            softmax_scale=softmax_scale,182            causal=causal,183            **flash_kwargs,184        )185        attn_output = attn_output.view(batch_size, -1, attn_output.size(-2), attn_output.size(-1))186 187    else:188        # print(f"scale {softmax_scale}")189        attn_output = flash_attn_func(190            query_states,191            key_states,192            value_states,193            # dropout,194            softmax_scale=softmax_scale,195            causal=causal,196            **flash_kwargs,197        )198 199    return attn_output200 201 202class FlashAttentionKwargs(TypedDict, total=False):203    cu_seq_lens_q: Optional[torch.LongTensor]204    cu_seq_lens_k: Optional[torch.LongTensor]205    max_length_q: Optional[int]206    max_length_k: Optional[int]207 208 209# _use_top_left_mask = flash_attn_supports_top_left_mask()210 211_use_top_left_mask = False212 213 214def flash_attention_forward(215    module: torch.nn.Module,216    query: torch.Tensor,217    key: torch.Tensor,218    value: torch.Tensor,219    attention_mask: Optional[torch.Tensor],220    dropout: float = 0.0,221    scaling: Optional[float] = None,222    sliding_window: Optional[int] = None,223    softcap: Optional[float] = None,224    **kwargs,225) -> Tuple[torch.Tensor, None]:226    # This is before the transpose227    seq_len = query.shape[1]228 229    # FA2 uses non-transposed inputs230    query = query.transpose(1, 2)231    key = key.transpose(1, 2)232    value = value.transpose(1, 2)233 234    # In PEFT, usually we cast the layer norms in float32 for training stability reasons235    # therefore the input hidden states gets silently casted in float32. Hence, we need236    # cast them back in the correct dtype just to be sure everything works as expected.237    # This might slowdown training & inference so it is recommended to not cast the LayerNorms238    # in fp32. (usually our RMSNorm modules handle it correctly)239    target_dtype = None240    if query.dtype == torch.float32:241        if torch.is_autocast_enabled():242            target_dtype = torch.get_autocast_gpu_dtype()243        # Handle the case where the model is quantized244        elif hasattr(module.config, "_pre_quantization_dtype"):245            target_dtype = module.config._pre_quantization_dtype246        else:247            target_dtype = next(layer for layer in module.modules() if isinstance(layer, torch.nn.Linear)).weight.dtype248 249    # FA2 always relies on the value set in the module, so remove it if present in kwargs to avoid passing it twice250    kwargs.pop("is_causal", None)251 252    attn_output = _flash_attention_forward(253        query,254        key,255        value,256        attention_mask,257        query_length=seq_len,258        is_causal=module.is_causal,259        dropout=dropout,260        softmax_scale=scaling,261        sliding_window=sliding_window,262        softcap=softcap,263        use_top_left_mask=_use_top_left_mask,264        target_dtype=target_dtype,265        **kwargs,266    )267 268    return attn_output, None