[Perf] Deepseekv3 performance optimization for eager mode

[ganyi1996ppo]

What this PR does / why we need it?

Deepseek v3 now adopt vanilla chunked prefill on MLA part which is ineffcient for computing but necessary for chunked prefill. Since PR #543 bring v0 scheduler into vllm-ascend, we can now adopt torch_npu._npu_flash_attention inside the mla backend for more performance boost. Also there are some redundant computation inside the rope, which is also removed. This PR should bring some performance gain for deepseek eager mode inference.

Does this PR introduce any user-facing change?

How was this patch tested?

[ganyi1996ppo]

Testing this PR on [2048in, 128out] * 16 scen with tp=2, ep=2 on deepseekv3-lite eager mode, got roughly 50% perf boost.

[ganyi1996ppo]

Found custom rotary_embedding was not opened before, after enable custom rotary_embedding, performance gain reach to roughly 100% compared to the main branch

[ganyi1996ppo]

BTW, this PR will enable v0 scheduler plug into engine v1 by default. Maybe we should discuss about this @wangxiyuan @wuhuikx .

[ganyi1996ppo]

@wuhuikx @wangxiyuan please help to review this pr

[ganyi1996ppo]

Disabled default v0 scheduler

@wuhuikx @wangxiyuan please help to review this pr

[wangxiyuan]

For this kind of change which override vllm. We should add note to describe why. L269 is the same.

[wuhuikx]

the DeepseekScalingRotaryEmbeeding is not a custom op in stock vLLM? Can we file a PR to vLLM to make it a custom op, after that can use the forward_oot, right?

[ganyi1996ppo]

it is a custom op, but there are some thing exist makes us needs for this changes.

• native deepseek rope cache the cos_sin differently compared with the naive implement in deepseek huggingface, which we found is more ascend friendly compared with vllm's impl
• its just override the forward rather than reuse its custom op's interface.

I'll add some comments on those line to explain the specific reason

[wuhuikx]

query_lens is a device tensor? if so, many D2H here, is this operation necessary?

[ganyi1996ppo]

query_lens is actually a cpu tensor, so no d2h operation will happened here, you can refer to line 220

[wuhuikx]

In prefill, we use MHA for computation, then the head_dim = nope_dim + rope_dim (192), while in decode, the absorbed and move_elision strategies are adopt, the head_dim=nope_dim, and we don't need pad, am I right?

[ganyi1996ppo]

You are definately right, this padding dim is used for prefill to padding the tensor. Not just for v_head_dim vs (qk_rope + qk_nope), but also for the 128 divisble head_dim alignment requirements for the _npu_flash_attention

[ganyi1996ppo]

The above data is measured under the v0 scheduler + v1 engine, if v0 scheduler is not enable, vanilla chunked prefill will be used to guarantee the functionality, which is combined with bunch of small ops. Vanilla operation harm the performance enoumously with lots of additional host and device operation, so we strongly recommand to adopt scheduler v0 for perf test by just add one single line.

To enable v0 scheduler in engine v1, you can pass additional_config with dict ascend_scheduler_config. For example, in your LLM offline_inference script, you should pass additional_config with kwarg ascend_scheduler_config to activate v0 scheduler in v1 engine in LLM:

    llm = LLM(model="/data/weights/deepseek-ai/deepseekv3-lite-base-latest",
              tensor_parallel_size=2,
              enforce_eager=True,
              trust_remote_code=True,
              max_model_len=1024,
              additional_config={"ascend_scheduler_config": {}})

[wangxiyuan]

use logger from vllm is better. And line 39 can be moved to L33. We can fix the nit later.

[wangxiyuan]

The above data is measured under the v0 scheduler + v1 engine, if v0 scheduler is not enable, vanilla chunked prefill will be used to guarantee the functionality, which is combined with bunch of small ops. Vanilla operation harm the performance enoumously with lots of additional host and device operation, so we strongly recommand to adopt scheduler v0 for perf test by just add one single line.

To enable v0 scheduler in engine v1, you can pass additional_config with dict ascend_scheduler_config. For example, in your LLM offline_inference script, you should pass additional_config with kwarg ascend_scheduler_config to activate v0 scheduler in v1 engine in LLM:

    llm = LLM(model="/data/weights/deepseek-ai/deepseekv3-lite-base-latest",
              tensor_parallel_size=2,
              enforce_eager=True,
              trust_remote_code=True,
              max_model_len=1024,
              additional_config={"ascend_scheduler_config": {}})

This info is very useful. Consider we plan to write a feature doc about ascend scheduler, this info can be involved in.