Implementation:OpenGVLab InternVL LLaMA2 Flash Attention Patch

Knowledge Sources	OpenGVLab_InternVL
Domains	Flash Attention, Monkey Patching, LLaMA, Performance
Last Updated	2026-02-07 14:00 GMT

Overview

Monkey-patches LLaMA 2 attention to use Flash Attention for faster and more memory-efficient self-attention computation, originally adapted from the FastChat project.

Description

This module provides a complete replacement for LlamaAttention.forward and LlamaModel._prepare_decoder_attention_mask to enable Flash Attention in LLaMA 2 models.

The custom forward function:

• Projects hidden states into Q, K, V tensors with shape (bsz, seq_len, num_heads, head_dim).
• Supports Grouped Query Attention (GQA) by reading num_key_value_heads from the attention module.
• Applies rotary position embeddings via a custom apply_rotary_pos_emb that gathers cos/sin values by position IDs.
• For unmasked inputs: calls flash_attn_func with causal=True for efficient full-sequence attention.
• For masked inputs: uses unpad_input to remove padding, packs K/V together, and calls flash_attn_varlen_kvpacked_func for variable-length sequences, then restores padding with pad_input.
• Supports past_key_value for incremental decoding (requires flash-attn >= 2.1.0).

The custom _prepare_decoder_attention_mask:

• Passes through a boolean key-padding mask instead of constructing the standard causal float attention mask.
• Returns None when all tokens are unmasked (training with full samples) for faster execution.

replace_llama2_attn_with_flash_attn performs the monkey-patching after checking GPU capability (warns if below A100/H100).

A test function validates correctness against the original FastChat implementation and tests incremental past-key-value decoding.

Usage

Call replace_llama2_attn_with_flash_attn() before loading a LLaMA 2 model to enable Flash Attention. This is typically done at the start of training or inference scripts for LLaMA 2-based InternVL models.

Code Reference

Source Location

• Repository: OpenGVLab_InternVL
• File: internvl_chat/internvl/patch/llama2_flash_attn_monkey_patch.py
• Lines: 1-238

Signature

def apply_rotary_pos_emb(q, k, cos_sin, position_ids): ...

def forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.Tensor] = None,
    past_key_value: Optional[Tuple[torch.Tensor]] = None,
    output_attentions: bool = False,
    use_cache: bool = False,
    padding_mask: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: ...

def _prepare_decoder_attention_mask(
    self, attention_mask, input_shape, inputs_embeds, past_key_values_length
): ...

def replace_llama2_attn_with_flash_attn(): ...

Import

from internvl.patch.llama2_flash_attn_monkey_patch import replace_llama2_attn_with_flash_attn

I/O Contract

Inputs

Name	Type	Required	Description
hidden_states	torch.Tensor	Yes	Input tensor of shape (batch_size, seq_len, hidden_size)
attention_mask	torch.Tensor	No	Boolean key-padding mask of shape (batch_size, seq_len)
position_ids	torch.Tensor	No	Position indices for rotary embeddings
past_key_value	Tuple[torch.Tensor]	No	Cached key/value tensors for incremental decoding

Outputs

Name	Type	Description
attn_output	torch.Tensor	Attention output of shape (batch_size, seq_len, hidden_size)
attn_weights	None	Always None (flash attention does not return attention weights)
past_key_value	Optional[Tuple]	Updated key/value cache if use_cache=True

Usage Examples

Basic Usage

from internvl.patch.llama2_flash_attn_monkey_patch import replace_llama2_attn_with_flash_attn

# Patch LLaMA 2 before loading model
replace_llama2_attn_with_flash_attn()

# Now load and use LLaMA 2 model normally
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b")

Related Pages

• Principle:OpenGVLab_InternVL_Flash_Attention_Patching