Support multistream of MLA vector operations #1135
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
sdmyzlp
changed the title
github-actions
bot
added
documentation
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
sdmyzlp
force-pushed
the
br_multi_stream_mla
branch
2 times, most recently
from
546e6c7 to
5b20d64
Compare
sdmyzlp
force-pushed
the
br_multi_stream_mla
branch
2 times, most recently
from
d4eac5b to
4967ade
Compare
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
sdmyzlp
force-pushed
the
br_multi_stream_mla
branch
from
4967ade to
be59b9b
Compare
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
sdmyzlp
force-pushed
the
br_multi_stream_mla
branch
from
be59b9b to
dcea3aa
Compare
sdmyzlp
force-pushed
the
br_multi_stream_mla
branch
from
dcea3aa to
980b356
Compare
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
1 similar comment
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
sdmyzlp
force-pushed
the
br_multi_stream_mla
branch
from
980b356 to
b7e7700
Compare
sdmyzlp
force-pushed
the
br_multi_stream_mla
branch
from
b7e7700 to
fd0f6fa
Compare
wangxiyuan
approved these changes
Jun 11, 2025
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
With the expected overlaping being:
```
| q_rmsnorm | | kv_norm_rope_cache | | q_rope |
| matmul W_DQ | matmul W_DKV | matmul W_UQ | split | matmul W_KV_T |
```
Controlled by `torchair_graph_config.enable_multistream_mla`, defaulted
to False.
Signed-off-by: sdmyzlp <lrwei2@petalmail.com>
sdmyzlp
force-pushed
the
br_multi_stream_mla
branch
from
fd0f6fa to
c89bf7b
Compare
|
there is a pr for refactor. I suggest to block this PR until that one is merged #1169 |
| decode_k_pe, decode_k_nope = self.exec_kv( | ||
| hidden_states_or_kv_c_normed, cos, sin, kv_cache, | ||
| attn_metadata.slot_mapping) | ||
| with npu_stream_switch("mla_secondary", |
There was a problem hiding this comment.
So, this npu_stream_switch can be used inside the torchair right? if that's so, can we further optimize the dbo path of deepseek to gain more performance boost.
There was a problem hiding this comment.
Discussed with torchair colleagues, the with torch.npu.stream(...): scheme is not supported by torchair graph mode while npu_stream_switch only works under graph mode, so I recommend leaving this to later PRs:
- to provide a joint version of
npu_stream_switchencapsulating both cases, as well as, - adding graph mode support for dbo
| hidden_states_or_q_c = self.q_a_layernorm(ckq) | ||
| use_multistream_mla = (self.enable_multistream_mla | ||
| and attn_metadata is not None | ||
| and attn_metadata.num_decodes > 0) |
There was a problem hiding this comment.
So, multistream mla will not be triggered if there are only prefill requests?
There was a problem hiding this comment.
Yes, current multistream utilities only support graph mode, i.e. npu_stream_switch returns contextlib.nullcontext() on non-graph mode, which is the case for prefill.
ganyi1996ppo
approved these changes
Jun 12, 2025
wangxiyuan
added a commit
to wangxiyuan/vllm-ascend
that referenced
this pull request
Jun 12, 2025
…)" This reverts commit e72f94e.
wangxiyuan
added a commit
to wangxiyuan/vllm-ascend
that referenced
this pull request
Jun 13, 2025
…)" This reverts commit e72f94e. Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>
momo609
pushed a commit
to momo609/vllm-ascend
that referenced
this pull request
Jun 17, 2025
### What this PR does / why we need it?
Move all vector operations to a secondary stream, with the expected
overlaping being:
```
| q_rmsnorm | | kv_norm_rope_cache | | q_rope |
| matmul W_DQ | matmul W_DKV | index | index | matmul W_UQ | split | matmul W_KV_T |
```
Currently, the `IndexByTensor` operators introduced by computation of
`cos` and `sin` can't be offloaded to the secondary stream due to a
known bug of graph fusion optimization pass. So we instead keep it in
the main stream, only requires it be computed before `matmul W_UQ` to
avoid hindering later overlapping. The problem may be solved by later
optimization (vllm-project#993), which hoists the computation of `cos` and `sin` up
to the first layer.
### Does this PR introduce _any_ user-facing change?
Controlled by `torchair_graph_config.enable_multistream_mla`, defaulted
to False.
### How was this patch tested?
Tested on 1x16 910 node, with tailored 2 layer DSKv2.
Signed-off-by: sdmyzlp <lrwei2@petalmail.com>
momo609
pushed a commit
to momo609/vllm-ascend
that referenced
this pull request
Jun 17, 2025
### What this PR does / why we need it?
Move all vector operations to a secondary stream, with the expected
overlaping being:
```
| q_rmsnorm | | kv_norm_rope_cache | | q_rope |
| matmul W_DQ | matmul W_DKV | index | index | matmul W_UQ | split | matmul W_KV_T |
```
Currently, the `IndexByTensor` operators introduced by computation of
`cos` and `sin` can't be offloaded to the secondary stream due to a
known bug of graph fusion optimization pass. So we instead keep it in
the main stream, only requires it be computed before `matmul W_UQ` to
avoid hindering later overlapping. The problem may be solved by later
optimization (vllm-project#993), which hoists the computation of `cos` and `sin` up
to the first layer.
### Does this PR introduce _any_ user-facing change?
Controlled by `torchair_graph_config.enable_multistream_mla`, defaulted
to False.
### How was this patch tested?
Tested on 1x16 910 node, with tailored 2 layer DSKv2.
Signed-off-by: sdmyzlp <lrwei2@petalmail.com>
Signed-off-by: wangxiaoxin (A) <wangxiaoxin7@huawei.com>
momo609
pushed a commit
to momo609/vllm-ascend
that referenced
this pull request
Jun 17, 2025
### What this PR does / why we need it?
Move all vector operations to a secondary stream, with the expected
overlaping being:
```
| q_rmsnorm | | kv_norm_rope_cache | | q_rope |
| matmul W_DQ | matmul W_DKV | index | index | matmul W_UQ | split | matmul W_KV_T |
```
Currently, the `IndexByTensor` operators introduced by computation of
`cos` and `sin` can't be offloaded to the secondary stream due to a
known bug of graph fusion optimization pass. So we instead keep it in
the main stream, only requires it be computed before `matmul W_UQ` to
avoid hindering later overlapping. The problem may be solved by later
optimization (vllm-project#993), which hoists the computation of `cos` and `sin` up
to the first layer.
### Does this PR introduce _any_ user-facing change?
Controlled by `torchair_graph_config.enable_multistream_mla`, defaulted
to False.
### How was this patch tested?
Tested on 1x16 910 node, with tailored 2 layer DSKv2.
Signed-off-by: sdmyzlp <lrwei2@petalmail.com>
Signed-off-by: wangxiaoxin (A) <wangxiaoxin7@huawei.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
What this PR does / why we need it?
Move all vector operations to a secondary stream, with the expected overlaping being:
Currently, the
IndexByTensoroperators introduced by computation ofcosandsincan't be offloaded to the secondary stream due to a known bug of graph fusion optimization pass. So we instead keep it in the main stream, only requires it be computed beforematmul W_UQto avoid hindering later overlapping. The problem may be solved by later optimization (#993), which hoists the computation ofcosandsinup to the first layer.Does this PR introduce any user-facing change?
Controlled by
torchair_graph_config.enable_multistream_mla, defaulted to False.How was this patch tested?
Tested on 1x16 910 node, with tailored 2 layer DSKv2.