revert torchair refactor

Signed-off-by: MengqingCao <[email protected]>

Author: MengqingCao

tests/ut/test_platform.py

@@ -1,12 +1,8 @@
 import importlib
-import unittest
-from datetime import timedelta
 from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
-from torch.distributed import ProcessGroup
-from torch.distributed.distributed_c10d import PrefixStore
 from vllm.config.compilation import CUDAGraphMode
 from vllm.platforms import PlatformEnum
 

vllm_ascend/compilation/acl_graph.py

@@ -40,7 +40,7 @@ class ACLGraphWrapper:
 
     The workflow of this wrapper in the aclgraph dispatching is as follows:
     1. At initialization, a runtime mode is assigned to the wrapper (FULL or
-    VLLM_COMPILE).
+    PIECEWISE).
     2. At runtime, the wrapper receives a runtime_mode and a
     batch_descriptor(key) from the forward context and blindly trust them
     for aclgraph dispatching.
@@ -126,7 +126,7 @@ def __call__(self, *args, **kwargs):
                 # Since we capture aclgraph for many different shapes and
                 # capturing is fast, we don't need to log it for every
                 # shape. E.g. we only log it for the first subgraph in
-                # VLLM_COMPILE mode.
+                # piecewise mode.
                 logger.debug("Capturing a aclgraph on (%s,%s)",
                              self.runtime_mode.name, entry.batch_descriptor)
             # validate that aclgraph capturing is legal at this point.
@@ -140,7 +140,7 @@ def __call__(self, *args, **kwargs):
 
             with ExitStack() as stack:
                 if self.aclgraph_options.gc_disable:
-                    # during every model forward for VLLM_COMPILE aclgraph
+                    # during every model forward for piecewise aclgraph
                     # mode, we will capture many pieces of aclgraphs
                     # (roughly one per layer). running gc again and again
                     # across layers will make the aclgraph capture very slow.
@@ -159,7 +159,7 @@ def __call__(self, *args, **kwargs):
                         # by converting it to weak ref,
                         # the original `output` will immediately be released
                         # to save memory. It is only safe to do this for
-                        # the last graph in VLLM_COMPILE aclgraph mode, because
+                        # the last graph in piecewise aclgraph mode, because
                         # the output of the last graph will not be used by
                         # any other acl graph.
                         output = weak_ref_tensors(output)

vllm_ascend/torchair/models/torchair_deepseek_v2.py

@@ -49,7 +49,6 @@
                                                RowParallelLinear,
                                                UnquantizedLinearMethod)
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.layers.mla import MLAModules
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -70,7 +69,7 @@
 
 from vllm_ascend import envs
 from vllm_ascend.ascend_config import get_ascend_config
-from vllm_ascend.models.layers.sfa import AscendSFAModules, Indexer
+from vllm_ascend.models.layers.sfa import Indexer
 from vllm_ascend.ops.weight_prefetch import maybe_npu_prefetch
 from vllm_ascend.quantization.quant_config import AscendLinearMethod
 from vllm_ascend.torchair.ops.torchair_fused_moe import TorchairAscendFusedMoE
@@ -79,9 +78,9 @@
 from vllm_ascend.utils import dispose_tensor, oproj_tp_enable, vllm_version_is
 
 if vllm_version_is("0.11.0"):
-    from vllm.model_executor.layers.mla import MultiHeadLatentAttention
+    from vllm.attention import Attention
 else:
-    from vllm.model_executor.layers.mla import MultiHeadLatentAttentionWrapper
+    from vllm.attention.layer import MLAAttention
 
 
 class TorchairDeepseekV2SiluAndMul(SiluAndMul):
@@ -486,6 +485,11 @@ def __init__(
         self.enable_shared_expert_dp = ascend_config.enable_shared_expert_dp
 
         if self.q_lora_rank is not None:
+            self.q_a_proj = ReplicatedLinear(self.hidden_size,
+                                             self.q_lora_rank,
+                                             bias=False,
+                                             quant_config=quant_config,
+                                             prefix=f"{prefix}.q_a_proj")
             self.q_a_layernorm = RMSNorm(self.q_lora_rank,
                                          eps=config.rms_norm_eps)
             self.q_b_proj = ColumnParallelLinear(q_lora_rank,
@@ -501,6 +505,12 @@ def __init__(
                                                bias=False,
                                                quant_config=quant_config,
                                                prefix=f"{prefix}.q_proj")
+        self.kv_a_proj_with_mqa = ReplicatedLinear(
+            self.hidden_size,
+            self.kv_lora_rank + self.qk_rope_head_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.kv_a_proj_with_mqa")
 
         self.kv_a_layernorm = RMSNorm(self.kv_lora_rank,
                                       eps=config.rms_norm_eps)
@@ -536,24 +546,6 @@ def __init__(
                 quant_config=quant_config,
                 prefix=f"{prefix}.o_proj")
 
-        print(30 * "=", f"q_lora_rank: {q_lora_rank}")
-        if self.q_lora_rank is not None:
-            self.fused_qkv_a_proj = MergedColumnParallelLinear(
-                self.hidden_size,
-                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
-                bias=False,
-                quant_config=quant_config,
-                prefix=f"{prefix}.fused_qkv_a_proj",
-                disable_tp=True)
-            self.kv_a_proj_with_mqa = None
-        else:
-            self.kv_a_proj_with_mqa = ReplicatedLinear(
-                self.hidden_size,
-                self.kv_lora_rank + self.qk_rope_head_dim,
-                bias=False,
-                quant_config=quant_config,
-                prefix=f"{prefix}.kv_a_proj_with_mqa")
-
         if rope_scaling:
             rope_scaling["rope_type"] = 'deepseek_yarn'
         self.rotary_emb = get_rope(qk_rope_head_dim,
@@ -568,59 +560,61 @@ def __init__(
             mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
             self.scaling = self.scaling * mscale * mscale
 
-        mla_modules = MLAModules(
-            kv_a_layernorm=self.kv_a_layernorm,
-            kv_b_proj=self.kv_b_proj,
-            rotary_emb=self.rotary_emb,
-            o_proj=self.o_proj,
-            fused_qkv_a_proj=self.fused_qkv_a_proj
-            if self.q_lora_rank is not None else None,
-            kv_a_proj_with_mqa=self.kv_a_proj_with_mqa
-            if self.q_lora_rank is None else None,
-            q_a_layernorm=self.q_a_layernorm
-            if self.q_lora_rank is not None else None,
-            q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,
-            q_proj=self.q_proj if self.q_lora_rank is None else None,
-            indexer=None,
-            is_sparse=hasattr(config, "index_topk"),
-            topk_indices_buffer=None,
-        )
         # In the MLA backend, kv_cache includes both k_c and
         # pe (i.e. decoupled position embeddings). In particular,
         # the concat_and_cache_mla op requires
         #     k_c.size(1) + k_pe.size(1) == kv_cache.size(2)
         # i.e.
         #     kv_lora_rank + qk_rope_head_dim == head_size
-
         if vllm_version_is("0.11.0"):
-            self.mla_attn = MultiHeadLatentAttention(
-                hidden_size=self.hidden_size,
+            self.mla_attn = Attention(
                 num_heads=self.num_local_heads,
+                head_size=self.kv_lora_rank + self.qk_rope_head_dim,
                 scale=self.scaling,
+                num_kv_heads=1,
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.attn",
+                use_mla=True,
+                # MLA Args
+                q_lora_rank=self.q_lora_rank,
+                kv_lora_rank=self.kv_lora_rank,
                 qk_nope_head_dim=self.qk_nope_head_dim,
                 qk_rope_head_dim=self.qk_rope_head_dim,
+                qk_head_dim=self.qk_head_dim,
                 v_head_dim=self.v_head_dim,
-                q_lora_rank=self.q_lora_rank,
-                kv_lora_rank=self.kv_lora_rank,
-                mla_modules=mla_modules,
-                cache_config=cache_config,
-                quant_config=quant_config,
-                prefix=prefix,
+                rotary_emb=self.rotary_emb,
+                q_proj=self.q_proj if self.q_lora_rank is None else None,
+                q_b_proj=self.q_b_proj
+                if self.q_lora_rank is not None else None,
+                kv_a_proj_with_mqa=self.kv_a_proj_with_mqa,
+                kv_a_layernorm=self.kv_a_layernorm,
+                kv_b_proj=self.kv_b_proj,
+                o_proj=self.o_proj,
             )
         else:
-            self.mla_attn = MultiHeadLatentAttentionWrapper(
-                hidden_size=self.kv_lora_rank + self.qk_rope_head_dim,
+            self.mla_attn = MLAAttention(
                 num_heads=self.num_local_heads,
                 scale=self.scaling,
                 qk_nope_head_dim=self.qk_nope_head_dim,
                 qk_rope_head_dim=self.qk_rope_head_dim,
                 v_head_dim=self.v_head_dim,
                 q_lora_rank=self.q_lora_rank,
                 kv_lora_rank=self.kv_lora_rank,
-                mla_modules=mla_modules,
                 cache_config=cache_config,
                 quant_config=quant_config,
-                prefix=prefix,
+                prefix=f"{prefix}.attn",
+                use_sparse=False,
+                indexer=None,
+                # MLA Args
+                rotary_emb=self.rotary_emb,
+                q_proj=self.q_proj if self.q_lora_rank is None else None,
+                q_b_proj=self.q_b_proj
+                if self.q_lora_rank is not None else None,
+                kv_a_proj_with_mqa=self.kv_a_proj_with_mqa,
+                kv_a_layernorm=self.kv_a_layernorm,
+                kv_b_proj=self.kv_b_proj,
+                o_proj=self.o_proj,
             )
 
     def forward(
@@ -658,9 +652,11 @@ def forward(
                                  dtype=hidden_states_or_q_c.dtype,
                                  device=hidden_states_or_q_c.device)
             forward_kwargs['output'] = output
-            output = self.mla_attn.mla_attn.impl.forward(
-                self.mla_attn, hidden_states_or_q_c, hidden_states, None,
-                kv_cache, attn_metadata, **forward_kwargs)
+            output = self.mla_attn.impl.forward(self.mla_attn,
+                                                hidden_states_or_q_c,
+                                                hidden_states, None, kv_cache,
+                                                attn_metadata,
+                                                **forward_kwargs)
             output = output.view(-1, output_shape[-1])
             return output
         else:
@@ -834,51 +830,55 @@ def __init__(
             prefix=f"{prefix}.indexer",
         )
 
-        sfa_modules = AscendSFAModules(
-            q_a_layernorm=self.q_a_layernorm
-            if self.q_lora_rank is not None else None,
-            q_proj=self.q_proj if self.q_lora_rank is None else self.q_b_proj,
-            q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,
-            kv_a_proj_with_mqa=self.kv_a_proj_with_mqa,
-            fused_qkv_a_proj=self.fused_qkv_a_proj
-            if self.q_lora_rank is not None else None,
-            kv_a_layernorm=self.kv_a_layernorm,
-            kv_b_proj=self.kv_b_proj,
-            o_proj=self.o_proj,
-            rotary_emb=self.rotary_emb,
-            indexer=self.indexer,
-            is_sparse=hasattr(config, "index_topk"),
-            topk_indices_buffer=None)
-
         if vllm_version_is("0.11.0"):
-            self.sfa_attn = MultiHeadLatentAttention(
-                hidden_size=self.hidden_size,
+            self.sfa_attn = Attention(
                 num_heads=self.num_local_heads,
+                head_size=self.kv_lora_rank + self.qk_rope_head_dim,
                 scale=self.scaling,
+                num_kv_heads=1,
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.attn",
+                use_mla=True,
+                # MLA Args
+                q_lora_rank=self.q_lora_rank,
+                kv_lora_rank=self.kv_lora_rank,
                 qk_nope_head_dim=self.qk_nope_head_dim,
                 qk_rope_head_dim=self.qk_rope_head_dim,
+                qk_head_dim=self.qk_head_dim,
                 v_head_dim=self.v_head_dim,
-                q_lora_rank=self.q_lora_rank,
-                kv_lora_rank=self.kv_lora_rank,
-                mla_modules=sfa_modules,
-                cache_config=cache_config,
-                quant_config=quant_config,
-                prefix=prefix,
+                rotary_emb=self.rotary_emb,
+                q_proj=self.q_proj if self.q_lora_rank is None else None,
+                q_b_proj=self.q_b_proj
+                if self.q_lora_rank is not None else None,
+                kv_a_proj_with_mqa=self.kv_a_proj_with_mqa,
+                kv_a_layernorm=self.kv_a_layernorm,
+                kv_b_proj=self.kv_b_proj,
+                o_proj=self.o_proj,
             )
         else:
-            self.sfa_attn = MultiHeadLatentAttentionWrapper(
-                hidden_size=self.hidden_size,
+            self.sfa_attn = MLAAttention(
                 num_heads=self.num_local_heads,
                 scale=self.scaling,
                 qk_nope_head_dim=self.qk_nope_head_dim,
                 qk_rope_head_dim=self.qk_rope_head_dim,
                 v_head_dim=self.v_head_dim,
                 q_lora_rank=self.q_lora_rank,
                 kv_lora_rank=self.kv_lora_rank,
-                mla_modules=sfa_modules,
                 cache_config=cache_config,
                 quant_config=quant_config,
-                prefix=prefix,
+                prefix=f"{prefix}.attn",
+                use_sparse=True,
+                indexer=self.indexer,
+                # MLA Args
+                rotary_emb=self.rotary_emb,
+                q_proj=self.q_proj if self.q_lora_rank is None else None,
+                q_b_proj=self.q_b_proj
+                if self.q_lora_rank is not None else None,
+                kv_a_proj_with_mqa=self.kv_a_proj_with_mqa,
+                kv_a_layernorm=self.kv_a_layernorm,
+                kv_b_proj=self.kv_b_proj,
+                o_proj=self.o_proj,
             )
 
     def forward(
@@ -917,9 +917,8 @@ def forward(
         output = torch.empty(output_shape,
                              dtype=hidden_states.dtype,
                              device=hidden_states.device)
-        self.sfa_attn.sfa_attn.impl.forward(hidden_states, kv_cache,
-                                            attn_metadata, need_gather_q_kv,
-                                            output)
+        self.sfa_attn.impl.forward(hidden_states, kv_cache, attn_metadata,
+                                   need_gather_q_kv, output)
         output = output.view(-1, output_shape[-1])
         return output
 

vllm_ascend/torchair/torchair_mla.py

@@ -664,7 +664,6 @@ def __init__(
         self.kv_a_layernorm = kwargs.get('kv_a_layernorm', None)
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
         self.tp_size = get_tensor_model_parallel_world_size()
-        self.fused_qkv_a_proj = kwargs.get('fused_qkv_a_proj', None)
 
         ascend_config = get_ascend_config()
         self.torchair_graph_enabled = ascend_config.torchair_graph_config.enabled
@@ -914,14 +913,7 @@ def exec_kv(
         B = hidden_states.shape[0]
         N = self.num_kv_heads
         S = 1
-        if self.fused_qkv_a_proj is not None:
-            qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
-            _, kv = qkv_lora.split(
-                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
-                dim=-1,
-            )
-        else:
-            kv = self.kv_a_proj_with_mqa(hidden_states)[0]
+        kv = self.kv_a_proj_with_mqa(hidden_states)[0]
         # npu_kv_rmsnorm_rope_cache needs [B, N, S, D]
         kv = kv.view(B, N, S, self.kv_lora_rank + self.qk_rope_head_dim)
         cache_mode = "PA_NZ" if self.enable_kv_nz else "PA"
@@ -950,14 +942,7 @@ def exec_kv_prefill(
         B = hidden_states.shape[0]
         N = self.num_kv_heads
         S = 1
-        if self.fused_qkv_a_proj is not None:
-            qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
-            _, kv = qkv_lora.split(
-                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
-                dim=-1,
-            )
-        else:
-            kv = self.kv_a_proj_with_mqa(hidden_states)[0]
+        kv = self.kv_a_proj_with_mqa(hidden_states)[0]
         # npu_kv_rmsnorm_rope_cache needs [B, N, S, D]
         kv = kv.view(B, N, S, self.kv_lora_rank + self.qk_rope_head_dim)
         cache_mode = "PA_NZ" if self.enable_kv_nz else "PA"
@@ -1120,23 +1105,9 @@ def forward(
         self.running_chunkprefilll_with_torchair = self.torchair_graph_enabled and attn_metadata.attn_state == AscendAttentionState.ChunkedPrefill
         num_actual_toks = attn_metadata.num_actual_tokens
         if k_pe is None and not self.running_in_graph:
-            if self.fused_qkv_a_proj is not None:
-                qkv_lora = self.fused_qkv_a_proj(
-                    hidden_states_or_kv_c_normed)[0]
-                _, kv = qkv_lora.split(
-                    [
-                        self.q_lora_rank,
-                        self.kv_lora_rank + self.qk_rope_head_dim
-                    ],
-                    dim=-1,
-                )
-                kv_c, k_pe = kv.split(
+            kv_c, k_pe = self.kv_a_proj_with_mqa(
+                hidden_states_or_kv_c_normed)[0].split(
                     [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
-            else:
-                kv = self.kv_a_proj_with_mqa(hidden_states_or_kv_c_normed)[0]
-                kv_c, k_pe = self.kv_a_proj_with_mqa(
-                    hidden_states_or_kv_c_normed)[0].split(
-                        [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
             kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())
         else:
             kv_c_normed = hidden_states_or_kv_c_normed