# monkeypatch.llama_attn_hijack_flash

`monkeypatch.llama_attn_hijack_flash`

Flash attention monkey patch for llama model

## Classes

Name | Description
--- | ---
FusedAttention | Fused QKV Attention layer for incrementally improved training efficiency
LlamaDecoderLayer | patched version of LlamaDecoderLayer to pass through the precalculated cu_seqlens

### FusedAttention

    monkeypatch.llama_attn_hijack_flash.FusedAttention(config, q, k, v, o)

Fused QKV Attention layer for incrementally improved training efficiency

### LlamaDecoderLayer

    monkeypatch.llama_attn_hijack_flash.LlamaDecoderLayer()

patched version of LlamaDecoderLayer to pass through the precalculated cu_seqlens

#### Methods

Name | Description
--- | ---
forward |

##### forward

    monkeypatch.llama_attn_hijack_flash.LlamaDecoderLayer.forward(
        hidden_states,
        attention_mask=None,
        position_ids=None,
        past_key_value=None,
        output_attentions=False,
        use_cache=False,
        padding_mask=None,
        cu_seqlens=None,
        max_seqlen=None,
    )

###### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
hidden_states | `torch.FloatTensor` | input to the layer of shape `(batch, seq_len, embed_dim)` | required
attention_mask | `torch.FloatTensor`, optional | attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. | `None`
output_attentions | `bool`, optional | Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. | `False`
use_cache | `bool`, optional | If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). | `False`
past_key_value | `Tuple(torch.FloatTensor)`, optional | cached past key and value projection states | `None`

## Functions

Name | Description
--- | ---
flashattn_forward | Input shape: Batch x Time x Channel
flashattn_forward_with_s2attn | Input shape: Batch x Time x Channel
generate_qkv |

### flashattn_forward

    monkeypatch.llama_attn_hijack_flash.flashattn_forward(
        self,
        hidden_states,
        attention_mask=None,
        position_ids=None,
        past_key_value=None,
        output_attentions=False,
        use_cache=False,
        padding_mask=None,
        cu_seqlens=None,
        max_seqlen=None,
    )

Input shape: Batch x Time x Channel

attention_mask: [bsz, q_len]

### flashattn_forward_with_s2attn

    monkeypatch.llama_attn_hijack_flash.flashattn_forward_with_s2attn(
        self,
        hidden_states,
        attention_mask=None,
        position_ids=None,
        past_key_value=None,
        output_attentions=False,
        use_cache=False,
        padding_mask=None,
        cu_seqlens=None,
        max_seqlen=None,
    )

Input shape: Batch x Time x Channel

From: https://github.com/dvlab-research/LongLoRA/blob/main/llama_attn_replace.py

attention_mask: [bsz, q_len]

`cu_seqlens` will be ignored if provided `max_seqlen` will be ignored if provided

### generate_qkv

    monkeypatch.llama_attn_hijack_flash.generate_qkv(
        q,
        k,
        v,
        query_padding_mask=None,
        key_padding_mask=None,
        kvpacked=False,
        qkvpacked=False,
    )

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
q |  | (batch_size, seqlen_q, nheads, d) | required
k |  | (batch_size, seqlen_k, nheads_k, d) | required
v |  | (batch_size, seqlen_k, nheads_k, d) | required
query_padding_mask |  | (batch_size, seqlen), bool | `None`
key_padding_mask |  | (batch_size, seqlen), bool | `None`