feat: Attention Module Replacement transform passes (ADLS Group 9)

[Johnny1882]

Code submission for ADLS Group 9

Description

This pull request implements a comprehensive framework for attention mechanism optimization and several transform passes for transformer-based models. They focus on improving model speed and memory efficiency at inference time while preserving model quality.

Contributions

We implement an extendable transform framework for attention module replacement, and three independent transform passes:

1. Multi-Head Latent Attention (MLA)
2. Group Query Attention (GQA)
3. LoRA Linear Layer

The supported model of transform passes includes:

• GPT-2
• Llama

Important Notes

• Dependency Issues: GPT2 model includes GPT2SpdaAttention module (available in transformers v4.47.1 only) and MLA code have custom kernel implemented, requiring Triton. These dependencies are causing expected failures in the current PR checks.
• Transformers Version: To successfully run transform passes on the GPT2 model, please set your transformers version to 4.47.1:

Reference

• Multi-Head Latent Attention (MLA): DeepSeek V3
• GPT2 huggingface

[Aaron-Zhao123]

Cleaned in #288, so closing this one.