EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[Performance]: Custom AscendC Kernel of Multi-Step Prepare Input (#814)

Browse Source 
### What this PR does / why we need it?

- According to https://github.com/vllm-project/vllm-ascend/issues/807,
we pull request for customer ascendc kernel of multi-step.
- also a bug we found in multi_step_runner.py is fixed when we use
multi-step on V0 Engine.


### Does this PR introduce _any_ user-facing change?

no user-facing change


### How was this patch tested?
we add Unit Test file and offline inference file to test the custom
ascendc kernel. See test/ops/test_multi_step.py and
examples/offline_multi_step.py

---------

Signed-off-by: wan_danfeng <wonderful199082@126.com>
This commit is contained in:
Wan_Danfeng
2025-05-20 09:31:30 +08:00
committed by GitHub
parent 00e0243561
commit 5cf9ff18e9
11 changed files with 629 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
vllm_ascend/attention/attention_v1.py
View File

@@ -220,11 +220,11 @@ class AscendAttentionBackendImpl(AttentionImpl):
key: shape = [batch_size, seq_len, num_kv_heads * head_size]
value: shape = [batch_size, seq_len, num_kv_heads * head_size]
kv_cache: shape = [2, num_blocks, block_size,
num_kv_heads * head_size]
num_kv_heads, head_size]
key_cache = [num_blocks, block_size,
num_kv_heads * head_size]
num_kv_heads, head_size]
value_cache = [num_blocks, block_size,
num_kv_heads * head_size]
num_kv_heads, head_size]
attn_metadata: Metadata for attention.
Returns:
shape = [batch_size * seq_len, num_heads, head_size]