注意力机制 (attention/)

一、模块概述

attention目录包含针对不同模型特性优化的注意力后端，支持标准Transformer、MLA、SFA、DSA、FA3等多种架构。通过上下文并行（Context Parallel）支持超长序列推理。

graph LR
    A[AttentionBackend 基类] --> B[AscendAttentionBackend 标准]
    A --> C[AscendMLABackend MLA]
    A --> D[AscendSFABackend SFA]
    A --> E[AscendDSABackend DSA]
    A --> F[AscendFABackend FA3]
    
    B --> B1[Q/K/V矩阵乘法]
    B --> B2[Softmax归一化]
    
    C --> C1[低秩KV压缩]
    C --> C2[共享KV头]
    
    D --> D1[稀疏注意力模式]
    D --> D2[选择性计算]
    
    E --> E1[动态稀疏模式]
    E --> E2[索引器支持]
    
    F --> F1[分块计算]
    F --> F2[内存优化]
                

二、文件结构

attention/
├── __init__.py              # 入口：get_attn_backend_cls()
├── attention_v1.py          # 标准注意力后端
├── mla_v1.py                # MLA多头潜在注意力
├── sfa_v1.py                # SFA稀疏注意力
├── dsa_v1.py                # DSA动态稀疏注意力
├── fa3_v1.py                # Flash Attention 3
├── sparse_flash_mla.py      # 稀疏Flash MLA操作
├── attention_mask.py        # AttentionMaskBuilder
├── indexer.py               # AscendSFAIndexerBackend
├── dsa_attn_kv_plan.py      # DsaAttnKvPlan, get_dsa_attn_kv_plan()
├── sfa_kv_offload.py        # SFA KV缓存卸载
├── utils.py                 # PagedAttentionGraphParam, AscendDCPMetadata
├── context_parallel/        # 上下文并行
│   ├── __init__.py
│   ├── attention_cp.py     # CP基类
│   ├── common_cp.py        # DCPMetadataBuilderMixin, DCPImplMixin
│   ├── mla_cp.py           # AscendMlaDCPImpl
│   ├── sfa_cp.py           # AscendSFADCPImpl
│   └── dsa_cp.py           # DSA上下文并行

三、后端选择逻辑

graph TD
    A[get_attn_backend_cls] --> B{use_mla?}
    B -->|Yes| C{use_sparse?}
    B -->|No| D{use_compress?}
    C -->|Yes| E[AscendSFABackend]
    C -->|No| F[AscendMLABackend]
    D -->|Yes| G[AscendDSABackend]
    D -->|No| H[AscendAttentionBackend]
                

def get_attn_backend_cls(use_mla, use_sparse, use_compress):
    # 根据模型配置选择注意力后端
    if use_mla:
        if use_sparse:
            return "vllm_ascend.attention.sfa_v1.AscendSFABackend"
        else:
            return "vllm_ascend.attention.mla_v1.AscendMLABackend"
    elif use_compress:
        return "vllm_ascend.attention.dsa_v1.AscendDSABackend"
    else:
        return "vllm_ascend.attention.attention_v1.AscendAttentionBackend"

四、后端详解

五、标准注意力流程

sequenceDiagram
    participant W as Worker
    participant A as AscendAttentionBackend
    participant N as NPU
    
    W->>A: prepare.attn_metadata()
    A-->>W: AscendMetadata
    
    W->>A: forward_masked(hidden_states, metadata)
    A->>N: Q * K^T (矩阵乘法)
    N-->>A: 注意力分数
    A->>N: 应用注意力掩码
    A->>N: Softmax归一化
    A->>N: 注意力权重 * V
    N-->>A: 注意力输出
    A-->>W: 输出hidden_states
                

六、MLA注意力详解

MLA（Multi-head Latent Attention）是DeepSeek系列模型的核心创新，通过低秩KV压缩将注意力计算的内存占用减少75%。

sequenceDiagram
    participant U as 上游vLLM Worker
    participant M as AscendMLABackend
    participant O as ops/mla.py
    participant N as NPU
    
    U->>M: forward_batched(hidden_states, attn_metadata)
    M->>O: 调用mla_forward()
    O->>N: 低秩KV压缩
    O->>N: Q * K^T (压缩后)
    O->>N: Softmax
    O->>N: 注意力权重 * V
    N-->>O: 注意力输出
    O-->>M: 返回结果
    M-->>U: 输出hidden_states
                

class AscendMLABackend(AttentionBackend):
    def forward_batched(self, hidden_states, attn_metadata):
        # 调用ops/mla.py中的MLA算子
        return mla_forward(
            hidden_states,
            attn_metadata.kv_cache,
            attn_metadata.mla_metadata
        )

七、上下文并行 (context_parallel/)

上下文并行将长序列分割到多个设备上计算，支持超长序列推理。DCP（Disaggregated Context Parallel）将Prefill和Decode分离到不同设备。

context_parallel/
├── __init__.py
├── attention_cp.py     # 上下文并行基类
├── common_cp.py        # DCPMetadataBuilderMixin, DCPImplMixin
├── mla_cp.py           # AscendMlaDCPImpl, AscendMlaDCPMetadataBuilder
├── sfa_cp.py           # AscendSFADCPImpl, AscendSFADCPMetadataBuilder
└── dsa_cp.py           # DSA上下文并行