feat: enable torch_npu graph mode for Qwen-3 dense with TP support. #325
Conversation
Summary of Pending Items for this PRThis PR is a work in progress. The following items need to be completed before it's ready for final review:
|
liutongxuan
changed the title
liutongxuan
changed the title
yingxudeng
force-pushed
the
feat/qwen3_npu_native_main
branch
from
4549991 to
1a5e2f0
Compare
xllm/core/kernels/param.h
Outdated
| // for npu | ||
| torch::Tensor seq_lens; | ||
| int num_heads; | ||
| int num_kv_heads; |
There was a problem hiding this comment.
we can get num_heads and num_kv_heads form th shape of query and key.
There was a problem hiding this comment.
In the file xllm_2/xllm/core/layers/common/attention.cpp, you’ve pre-shaped the query and key parameters before passing them. However, different platforms may require slightly different shapes after the view operation. It might be more flexible to pass the original query, key along with num_heads_ and num_kv_heads_ parameters, and then perform the view operation inside the batch_prefill method. This approach would provide better cross-platform compatibility and clearer parameter handling.
yingxudeng
force-pushed
the
feat/qwen3_npu_native_main
branch
10 times, most recently
from
34afffc to
97f924a
Compare
…ilation structure.
yingxudeng
force-pushed
the
feat/qwen3_npu_native_main
branch
from
97f924a to
f37d69d
Compare
|
Could you please help review this PR when you have a moment? 🙏 @yq33victor @XuZhang99 |
| proto::CommUniqueIdList uids; | ||
| sync_master_node(master_node_addr, addr_info, uids); | ||
|
|
||
| CollectiveCommunicator comm(worker_global_rank, world_size, dp_size, ep_size); |
There was a problem hiding this comment.
todo:
增加类似npu_operator_backend 的参数，
在CollectiveCommunicator 初始化、create_process_groups 函数中区分是 atb 后端还是torch 后端的。
后续增加新模型时，根据 model_type自动选择 npu_operator_backend 是 ATB 还是 TORCH。
|
|
||
| CollectiveCommunicator comm(worker_global_rank, world_size, dp_size, ep_size); | ||
| const ParallelArgs* parallel_args = comm.parallel_args(); | ||
| #if defined(USE_MLU) || defined(USE_CUDA) |
There was a problem hiding this comment.
同上考虑，去掉 #if 这种，统一代码，在函数内部判断
| torch::Device device_; | ||
|
|
||
| protected: | ||
| #if defined(USE_NPU) |
There was a problem hiding this comment.
这块是因为 torch 的版本是 2.1.0 这个时候 backend 类里面还没有 shutdown() 方法，无法通过编译。如果要统一可能要升级 torch 版本，成本有些高，暂时这样
There was a problem hiding this comment.
use #if TORCH_VERSION_MAJOR >= 2 && TORCH_VERSION_MINOR >= 7
| // for npu | ||
| if (attn_mask.has_value()) { | ||
| attn_metadata.attn_mask = attn_mask.value(); | ||
| attn_metadata.seq_lens = params.kv_seq_lens.to(torch::kCPU); |
There was a problem hiding this comment.
这会导致开启 acl grpah 出问题，因为 attn_metadata.seq_lens 在 cpu 上面。
需要接入民超哥的新 attention
| activation_params.output = output; | ||
| activation_params.act_mode = hidden_act_; | ||
| activation_params.is_gated = is_gated_; | ||
| xllm::kernel::active(activation_params); |
There was a problem hiding this comment.
@yq33victor 大佬请教一下，类似这种 active 函数是不是可以都统一成空间申请在函数内，然后 output 再 return 出来呢？
这样用起来更加像 torch 一点，这块的代码也可以统一了。
| fused_layernorm_params.mode = kRmsNormMode; | ||
| fused_layernorm_params.eps = eps_; | ||
|
|
||
| xllm::kernel::fused_layernorm(fused_layernorm_params); |
2 participants
No description provided.