ascend 950 support qwen dense model

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

Author: wangyao-i

vllm_ascend/attention/attention_v1.py

@@ -50,7 +50,7 @@
                                                update_graph_params_workspaces)
 from vllm_ascend.ops.attention import vanilla_chunked_prefill
 from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_NZ, aligned_16, is_310p,
-                               nd_to_nz_2d, nd_to_nz_spec,
+                               is_Ascend950, nd_to_nz_2d, nd_to_nz_spec,
                                prefill_context_parallel_enable,
                                weak_ref_tensors)
 
@@ -703,15 +703,29 @@ def _forward_prefill_no_cache(
             mask = torch_npu.npu_format_cast(mask.contiguous(),
                                              ACL_FORMAT_FRACTAL_NZ)
 
-        torch_npu._npu_flash_attention(query=query,
-                                       key=key,
-                                       value=value,
-                                       mask=mask,
-                                       seq_len=attn_metadata.seq_lens,
-                                       scale_value=self.scale,
-                                       num_heads=self.num_heads,
-                                       num_kv_heads=self.num_kv_heads,
-                                       out=output)
+        if is_Ascend950():
+            num_tokens = attn_metadata.query_start_loc[-1]
+            output, _ = torch_npu.npu_fused_infer_attention_score_v2(
+                query[:num_tokens],
+                key[:num_tokens],
+                value[:num_tokens],
+                atten_mask=mask.to(torch.bool),
+                actual_seq_qlen=attn_metadata.query_lens.cumsum(0),
+                actual_seq_kvlen=attn_metadata.seq_lens.cumsum(0),
+                num_query_heads=self.num_heads,
+                num_key_value_heads=self.num_kv_heads,
+                input_layout="TND",
+                softmax_scale=self.scale)
+        else:
+            torch_npu._npu_flash_attention(query=query,
+                                           key=key,
+                                           value=value,
+                                           mask=mask,
+                                           seq_len=attn_metadata.seq_lens,
+                                           scale_value=self.scale,
+                                           num_heads=self.num_heads,
+                                           num_kv_heads=self.num_kv_heads,
+                                           out=output)
         assert output is not None
         return output[:num_tokens]
 
@@ -729,6 +743,27 @@ def _forward_prefill_cache_hit(
         block_table = attn_metadata.block_tables[:batch_size, :]
         num_block, block_size, _, _ = self.key_cache.shape  # type: ignore
 
+        if is_Ascend950():
+            compress_mask = compress_mask.to(torch.bool)
+            key = self.key_cache.transpose(1, 2)  # type: ignore
+            value = self.value_cache.transpose(1, 2)  # type: ignore
+            block_size = self.block_size
+            output, _ = torch_npu.npu_fused_infer_attention_score_v2(
+                query=query,
+                key=key,
+                value=value,
+                block_table=block_table,
+                atten_mask=compress_mask,
+                actual_seq_qlen=attn_metadata.query_lens.cumsum(0),
+                actual_seq_kvlen=attn_metadata.seq_lens,
+                num_query_heads=self.num_heads,
+                num_key_value_heads=self.num_kv_heads,
+                softmax_scale=self.scale,
+                spare_mode=2,
+                block_size=block_size,
+                input_layout="TND")
+            return output
+
         if block_size == 128:
             # TODO:The npu_fused_infer_attention_score op is planned to
             # be utilized in a wider range in upcoming versions.
@@ -777,18 +812,20 @@ def _forward_decode_only(
             # seq_lens_tensor needs to be transferred to the device for 310P.
             attn_metadata.seq_lens = \
                 attn_metadata.seq_lens.to(device=query.device)
+
+        batch_size = attn_metadata.seq_lens.shape[0]
+        block_size = 128
+        key = self.key_cache
+        value = self.value_cache
+        if self.key_cache is not None and self.value_cache is not None:
+            block_size = self.key_cache.shape[1]
+            key = self.key_cache.flatten(2, 3).contiguous()
+            value = self.value_cache.flatten(2, 3).contiguous()
+
         if self.sliding_window is not None and attn_metadata.seq_lens.shape[
                 0] == query.size(0):
-            batch_size = attn_metadata.seq_lens.shape[0]
-            block_size = 128
-            query = query.view(batch_size, 1, self.num_heads * self.head_size)
-            key = self.key_cache
-            value = self.value_cache
-            if self.key_cache is not None and self.value_cache is not None:
-                block_size = self.key_cache.shape[1]
-                key = self.key_cache.flatten(2, 3).contiguous()
-                value = self.value_cache.flatten(2, 3).contiguous()
-
+            query = query.view(batch_size, 1,
+                               self.num_heads * self.head_size)
             output, _ = torch_npu.npu_fused_infer_attention_score(
                 query,
                 key,
@@ -805,16 +842,33 @@ def _forward_decode_only(
 
             output = output.view(batch_size, self.num_heads, self.head_size)
         else:
-            torch_npu._npu_paged_attention(
-                query=query,
-                key_cache=self.key_cache,
-                value_cache=self.value_cache,
-                num_kv_heads=self.num_kv_heads,
-                num_heads=self.num_heads,
-                scale_value=self.scale,
-                block_table=attn_metadata.block_tables,
-                context_lens=attn_metadata.seq_lens,
-                out=output)
+            if is_Ascend950():
+                query = query.view(batch_size, 1,
+                                   self.num_heads * self.head_size)
+                output, _ = torch_npu.npu_fused_infer_attention_score_v2(
+                    query=query,
+                    key=key,
+                    value=value,
+                    actual_seq_kvlen=attn_metadata.seq_lens,
+                    num_query_heads=self.num_heads,
+                    num_key_value_heads=self.num_kv_heads,
+                    block_table=attn_metadata.block_tables[:batch_size],
+                    block_size=block_size,
+                    softmax_scale=self.scale,
+                    input_layout="BSH")
+                output = output.view(batch_size, self.num_heads,
+                                     self.head_size)
+            else:
+                torch_npu._npu_paged_attention(
+                    query=query,
+                    key_cache=self.key_cache,
+                    value_cache=self.value_cache,
+                    num_kv_heads=self.num_kv_heads,
+                    num_heads=self.num_heads,
+                    scale_value=self.scale,
+                    block_table=attn_metadata.block_tables,
+                    context_lens=attn_metadata.seq_lens,
+                    out=output)
         return output
 
     def _forward_v1_style(
@@ -862,7 +916,6 @@ def _forward_v1_style(
             num_block, block_size, -1)
         value = self.value_cache.view(  # type: ignore
             num_block, block_size, -1)
-
         output, _ = torch_npu.npu_fused_infer_attention_score(
             query=query,
             key=key,
@@ -1507,12 +1560,20 @@ def forward(
             if has_decode:
                 slot_mapping = attn_metadata.slot_mapping[:num_decode_tokens * self.pcp_size: self.pcp_size] \
                     if self.pcp_size * self.dcp_size > 1 else attn_metadata.slot_mapping[:num_decode_tokens]
-                torch_npu._npu_reshape_and_cache(
-                    key=key[:num_decode_tokens],
-                    value=value[:num_decode_tokens],
-                    key_cache=self.key_cache,
-                    value_cache=self.value_cache,
-                    slot_indices=slot_mapping)
+                if is_Ascend950():
+                    num_tokens = slot_mapping.shape[0]
+                    torch_npu.npu_scatter_pa_kv_cache(
+                        key=key[:num_tokens],
+                        value=value[:num_tokens].contiguous(),
+                        slot_mapping=slot_mapping,
+                        out=(self.key_cache, self.value_cache))
+                else:
+                    torch_npu._npu_reshape_and_cache(
+                        key=key[:num_decode_tokens],
+                        value=value[:num_decode_tokens],
+                        key_cache=self.key_cache,
+                        value_cache=self.value_cache,
+                        slot_indices=slot_mapping)
 
             if has_prefill:
                 if self.pcp_size > 1:
@@ -1526,18 +1587,27 @@ def forward(
                     key, value = all_kv.split([self.head_size, self.head_size],
                                               dim=-1)
 
-                torch_npu._npu_reshape_and_cache(
-                    key=key[self.pcp_size * num_decode_tokens:attn_metadata.
-                            num_actual_tokens_pcp_padded],
-                    value=value[self.pcp_size *
+                if is_Ascend950():
+                    num_tokens = attn_metadata.slot_mapping.shape[0]
+                    torch_npu.npu_scatter_pa_kv_cache(
+                        key=key[:num_tokens],
+                        value=value[:num_tokens].contiguous(),
+                        slot_mapping=attn_metadata.slot_mapping,
+                        out=(self.key_cache, self.value_cache))
+                else:
+                    torch_npu._npu_reshape_and_cache(
+                        key=key[self.pcp_size *
                                 num_decode_tokens:attn_metadata.
                                 num_actual_tokens_pcp_padded],
-                    key_cache=self.key_cache,
-                    value_cache=self.value_cache,
-                    slot_indices=attn_metadata.
-                    slot_mapping[self.pcp_size *
-                                 num_decode_tokens:attn_metadata.
-                                 num_actual_tokens_pcp_padded])
+                        value=value[self.pcp_size *
+                                    num_decode_tokens:attn_metadata.
+                                    num_actual_tokens_pcp_padded],
+                        key_cache=self.key_cache,
+                        value_cache=self.value_cache,
+                        slot_indices=attn_metadata.
+                        slot_mapping[self.pcp_size *
+                                     num_decode_tokens:attn_metadata.
+                                     num_actual_tokens_pcp_padded])
 
         forward_context: ForwardContext = get_forward_context()
         if not forward_context.capturing:

vllm_ascend/ops/rotary_embedding.py

@@ -26,7 +26,7 @@
     YaRNScalingRotaryEmbedding)
 
 from vllm_ascend.platform import NPUPlatform
-from vllm_ascend.utils import enable_custom_op, is_310p
+from vllm_ascend.utils import enable_custom_op, is_310p, is_Ascend950
 
 
 def _custom_rotary_embedding_enabled(query, neox_style, head_size):
@@ -405,7 +405,7 @@ def forward_oot(
         query: torch.Tensor,
         key: torch.Tensor,
     ):
-        if self.mrope_section != [16, 24, 24]:
+        if self.mrope_section != [16, 24, 24] or is_Ascend950():
             return super().forward_oot(positions, query, key)
 
         import torch_npu

vllm_ascend/sample/sampler.py

@@ -3,7 +3,7 @@
 from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler, random_sample
 from vllm.v1.sample.sampler import Sampler
 
-from vllm_ascend.utils import is_310p
+from vllm_ascend.utils import is_310p, is_Ascend950
 
 DEFAULT_LOGPROBS_MODE = "raw_logprobs"
 
@@ -25,8 +25,8 @@ def _apply_top_k_top_p(
         p: torch.Tensor,
     ) -> torch.Tensor:
         # npu_top_k_top_p uses the operator aclnnApplyTopKTopP, but aclnnApplyTopKTopP currently does not support 310P
-        if not is_310p() and p is not None and k is not None and 1 <= int(
-                k.max()) <= 1024:
+        if not is_310p() and not is_Ascend950() \
+            and p is not None and k is not None and 1 <= int(k.max()) <= 1024:
             # npu_top_k_top_p's parameter order is (logits, p, k), not (logits, k, p)
             return torch_npu.npu_top_k_top_p(logits, p, k)
 

vllm_ascend/utils.py

@@ -49,6 +49,7 @@
 
 _CUSTOM_OP_ENABLED = None
 _IS_310P = None
+_IS_ASCEND950 = None
 _SLEEP_MODE_ENABLED = None
 _CURRENT_STREAM = None
 _PREFETCH_STREAM = None
@@ -668,7 +669,8 @@ def register_ascend_customop(vllm_config: Optional[VllmConfig] = None):
 class AscendSocVersion(Enum):
     A2 = 0
     A3 = 1
-    UNDEFINED = 2
+    A5 = 2
+    UNDEFINED = 3
 
 
 _ascend_soc_version = None
@@ -681,6 +683,8 @@ def init_ascend_soc_version():
         _ascend_soc_version = AscendSocVersion.A2
     elif 250 <= soc_version <= 255:
         _ascend_soc_version = AscendSocVersion.A3
+    elif soc_version == 260:
+        _ascend_soc_version = AscendSocVersion.A5
     else:
         _ascend_soc_version = AscendSocVersion.UNDEFINED
 
@@ -945,3 +949,10 @@ def get_flashcomm2_reorgnized_batch_ids(global_tp_size) -> list[list[int]]:
         reorgnized_batch_ids.append(ranks)
 
     return reorgnized_batch_ids
+
+
+def is_Ascend950():
+    global _IS_ASCEND950
+    if _IS_ASCEND950 is None:
+        _IS_ASCEND950 = (get_ascend_soc_version() == AscendSocVersion.A5)
+    return _IS_ASCEND950

vllm_ascend/worker/worker_v1.py

@@ -47,8 +47,8 @@
 from vllm_ascend.device_allocator.camem import CaMemAllocator
 from vllm_ascend.distributed.parallel_state import init_ascend_model_parallel
 from vllm_ascend.platform import NPUPlatform
-from vllm_ascend.utils import (init_ascend_soc_version, is_enable_nz,
-                               prefill_context_parallel_enable,
+from vllm_ascend.utils import (init_ascend_soc_version, is_Ascend950,
+                               is_enable_nz, prefill_context_parallel_enable,
                                register_ascend_customop, sleep_mode_enabled,
                                try_register_lib, vllm_version_is)
 from vllm_ascend.worker.model_runner_v1 import NPUModelRunner
@@ -342,7 +342,8 @@ def compile_or_warm_up_model(self) -> None:
             self.model_runner.capture_model()
         # Call ATB matmul to warm up; otherwise, the first operation (ReshapeAndCache)
         # may cause performance degradation at runtime.
-        self._warm_up_atb()
+        if not is_Ascend950():
+            self._warm_up_atb()
         # Reset the seed to ensure that the random state is not affected by
         # the model initialization and profiling.
         NPUPlatform.seed_everything(self.model_config.seed)