feature for triton rerope

[xinSky00]

Purpose

This PR implements a Triton kernel for ReRoPE to enable efficient context length extension in vLLM. The ReRoPE-based approach delivers 3-5x speedup for long sequences.

It implements segment-wise attention computation, applying full rotary embeddings within a fixed window while constraining positional encodings beyond the window boundary, enabling models to handle sequences beyond their pre-training length without fine-tuning.

Modifications

1. New Environment Variable: VLLM_ATTENTION_BACKEND

• Must be set to TRITON_ATTN_VLLM_V1 to enable the Triton-based backend
• Usage: export VLLM_ATTENTION_BACKEND="TRITON_ATTN_VLLM_V1"
• Default: FLASH_ATTN (standard FlashAttention backend)

2. Modified Model Configuration Parameter: max_position_embeddings

• Users must adjust this parameter via --hf-overrides to match their target input length
• This ensures the RoPE embeddings are properly computed for extended sequences
• Example: --hf-overrides '{"max_position_embeddings": 327680}'

3. ReRoPE-specific parameters: rerope_window and training_length should be configured based on the model's original pre-training length

• These values determine the segment boundaries for attention computation and must align with the model's original training configuration

Test

run the file of offline_inference_rerope.py

• os.environ["VLLM_ATTENTION_BACKEND"] = "TRITON_ATTN_VLLM_V1"
• os.environ["VLLM_USE_REROPE"] = "true"
• model: Qwen2.5-14B-Instruct
• Dataset: multifieldqa_zh.jsonl
• results
  • prompt length: about 130k tokens
    
  • prompt length: about 315k tokens