From 2bce4902e3784caac5e690244c663826db26b904 Mon Sep 17 00:00:00 2001 From: Maojun Xu Date: Mon, 17 Aug 2026 10:51:51 +0800 Subject: [PATCH 1/3] perf: enable optimized qwen3 nz decode routing. - route eligible low-token decode shapes through the non-fused kernel\n- pin xllm-ops and catlass dependencies containing the optimized tiling --- third_party/xllm_atb_layers | 2 +- third_party/xllm_ops | 2 +- .../npu/npu_qwen3_decoder_layer_impl.cpp | 84 ++++++++++++++++++- .../layers/npu/npu_qwen3_decoder_layer_impl.h | 4 + 4 files changed, 87 insertions(+), 5 deletions(-) diff --git a/third_party/xllm_atb_layers b/third_party/xllm_atb_layers index 92171d597f..7603995e86 160000 --- a/third_party/xllm_atb_layers +++ b/third_party/xllm_atb_layers @@ -1 +1 @@ -Subproject commit 92171d597ff87d468e55fc8b0224d02b6cf65306 +Subproject commit 7603995e8676d08827ae6ceeee3482a5d4bf9ff9 diff --git a/third_party/xllm_ops b/third_party/xllm_ops index d4a441a60b..a1a6de04a5 160000 --- a/third_party/xllm_ops +++ b/third_party/xllm_ops @@ -1 +1 @@ -Subproject commit d4a441a60b8b2a42ae2e722e662b21d508d6781c +Subproject commit a1a6de04a5934ad2c79cc8e2c7d222e62cf21088 diff --git a/xllm/core/layers/npu/npu_qwen3_decoder_layer_impl.cpp b/xllm/core/layers/npu/npu_qwen3_decoder_layer_impl.cpp index ab240d0f95..465eb78997 100644 --- a/xllm/core/layers/npu/npu_qwen3_decoder_layer_impl.cpp +++ b/xllm/core/layers/npu/npu_qwen3_decoder_layer_impl.cpp @@ -28,6 +28,8 @@ limitations under the License. #include "core/framework/config/parallel_config.h" #include "core/framework/config/scheduler_config.h" #include "core/layers/npu/loader/qwen3_decoder_loader.h" +#include "operations/aclnn/ops/quant_matmul_nz_decode_operation.h" +#include "operations/aclnn/utils/utils.h" #include "operations/fusion/mlp/mlp.h" #include "util/rec_model_utils.h" @@ -39,6 +41,24 @@ namespace xllm { namespace layer { const uint64_t WEIGHT_COUNT_PER_LAYER = 56; +namespace { + +constexpr int64_t kOptimizedHiddenSize = 5120; +constexpr int64_t kOptimizedIntermediateSizePerRank = 3200; +constexpr int64_t kOptimizedQkvSizePerRank = 1280; +constexpr int64_t kMaxOptimizedDecodeTokens = 16; + +bool is_low_latency_decode_bucket(const torch::Tensor& input) { + if (!input.defined() || (input.dim() != 2 && input.dim() != 3) || + input.size(-1) != kOptimizedHiddenSize) { + return false; + } + const int64_t hidden_size = input.size(-1); + const int64_t token_count = input.numel() / hidden_size; + return token_count > 0 && token_count <= kMaxOptimizedDecodeTokens; +} + +} // namespace void NpuQwen3DecoderLayerImpl::param_from_args( atb_speed::qwen::QwenLayerParam& param, @@ -109,6 +129,9 @@ void NpuQwen3DecoderLayerImpl::param_from_args( } initialize_parallel_parameters(param, parallel_args); initialize_quantization_parameters(param); + param.enableQuantMatmulNzGateUpDecode = false; + param.enableQuantMatmulNzDownDecode = false; + param.enableQuantMatmulNzQkvDecode = false; if (isPrefill) { param.enableAclnnRmsNorm = @@ -197,6 +220,36 @@ NpuQwen3DecoderLayerImpl::NpuQwen3DecoderLayerImpl(const ModelContext& context) param_from_args(prefill_param_, model_args, parallel_args, true); param_from_args(decode_graph_param_, model_args, parallel_args, false); + decode_optimized_graph_param_ = decode_graph_param_; + const int64_t qkv_size_per_rank = + (decode_graph_param_.numAttentionHeadsPerRank + + 2 * decode_graph_param_.numKeyValueHeadsPerRank) * + decode_graph_param_.hiddenSizePerAttentionHead; + const bool is_target_hidden_size = + model_args.hidden_size() == kOptimizedHiddenSize; + const bool is_target_mlp_shape = + is_target_hidden_size && + model_args.intermediate_size() == + kOptimizedIntermediateSizePerRank * parallel_args.world_size(); + const bool is_target_qkv_shape = + is_target_hidden_size && qkv_size_per_rank == kOptimizedQkvSizePerRank; + const bool supports_low_latency_decode = + quantize_type_ == "w8a8" && decode_graph_param_.isBF16 && + (is_target_mlp_shape || is_target_qkv_shape) && + decode_graph_param_.enableAclGraphPagedAttention && + decode_graph_param_.matmulBackend == + atb_speed::common::OpBackend::ACLNN && + !decode_graph_param_.enableLora && !decode_graph_param_.enableFlashComm && + atb_speed::common::IsA2(); + const bool quant_matmul_nz_decode_available = + supports_low_latency_decode && + atb_speed::common::QuantMatmulNzDecodeOperation::is_available(); + decode_optimized_graph_param_.enableQuantMatmulNzGateUpDecode = + quant_matmul_nz_decode_available && is_target_mlp_shape; + decode_optimized_graph_param_.enableQuantMatmulNzDownDecode = + quant_matmul_nz_decode_available && is_target_mlp_shape; + decode_optimized_graph_param_.enableQuantMatmulNzQkvDecode = + quant_matmul_nz_decode_available && is_target_qkv_shape; decode_eager_param_ = decode_graph_param_; decode_eager_param_.enableAclGraphPagedAttention = false; atb_weight_tensors_.resize(WEIGHT_COUNT_PER_LAYER); @@ -227,7 +280,15 @@ int64_t NpuQwen3DecoderLayerImpl::init_layer() { if (quantize_type_ == "w8a8") { Qwen3DecoderLoader* qwen3_loader = dynamic_cast(loader_.get()); - if (qwen3_loader && qwen3_loader->down_proj_quantized()) { + const bool down_proj_quantized = + qwen3_loader != nullptr && qwen3_loader->down_proj_quantized(); + decode_optimized_graph_param_.enableQuantMatmulNzGateUpDecode = + decode_optimized_graph_param_.enableQuantMatmulNzGateUpDecode && + down_proj_quantized; + decode_optimized_graph_param_.enableQuantMatmulNzDownDecode = + decode_optimized_graph_param_.enableQuantMatmulNzDownDecode && + down_proj_quantized; + if (down_proj_quantized) { auto update_down_proj = [](atb_speed::qwen::QwenLayerParam& p) { p.linearDescs[atb_speed::common::DOWN_LINEAR_INDEX] = static_cast(LinearTypeV2::W8A8); @@ -237,6 +298,7 @@ int64_t NpuQwen3DecoderLayerImpl::init_layer() { }; update_down_proj(prefill_param_); update_down_proj(decode_graph_param_); + update_down_proj(decode_optimized_graph_param_); update_down_proj(decode_eager_param_); } if (qwen3_loader && !qwen3_loader->o_proj_quantized()) { @@ -251,13 +313,23 @@ int64_t NpuQwen3DecoderLayerImpl::init_layer() { }; update_o_proj(prefill_param_); update_o_proj(decode_graph_param_); + update_o_proj(decode_optimized_graph_param_); update_o_proj(decode_eager_param_); } } + enable_optimized_decode_graph_ = + decode_optimized_graph_param_.enableQuantMatmulNzGateUpDecode || + decode_optimized_graph_param_.enableQuantMatmulNzDownDecode || + decode_optimized_graph_param_.enableQuantMatmulNzQkvDecode; + CHECK_OPERATION_STATUS_RETURN(init_node(prefill_node_, prefill_param_)); CHECK_OPERATION_STATUS_RETURN( init_node(decode_graph_node_, decode_graph_param_)); + if (enable_optimized_decode_graph_) { + CHECK_OPERATION_STATUS_RETURN( + init_node(decode_optimized_graph_node_, decode_optimized_graph_param_)); + } CHECK_OPERATION_STATUS_RETURN( init_node(decode_eager_node_, decode_eager_param_)); @@ -325,8 +397,14 @@ torch::Tensor NpuQwen3DecoderLayerImpl::forward(torch::Tensor& x, const bool use_graph_decode_input = ::xllm::ExecutionConfig::get_instance().enable_graph() && input_params.graph.tiling_data.defined(); - auto& decode_node = - use_graph_decode_input ? decode_graph_node_ : decode_eager_node_; + const bool use_low_latency_decode = use_graph_decode_input && + enable_optimized_decode_graph_ && + is_low_latency_decode_bucket(x); + atb_speed::Model::Node& decode_node = + use_graph_decode_input + ? (use_low_latency_decode ? decode_optimized_graph_node_ + : decode_graph_node_) + : decode_eager_node_; build_node_variant_pack(decode_node, x, cos_pos, diff --git a/xllm/core/layers/npu/npu_qwen3_decoder_layer_impl.h b/xllm/core/layers/npu/npu_qwen3_decoder_layer_impl.h index fceb1452ab..004db69eb4 100644 --- a/xllm/core/layers/npu/npu_qwen3_decoder_layer_impl.h +++ b/xllm/core/layers/npu/npu_qwen3_decoder_layer_impl.h @@ -64,6 +64,7 @@ class NpuQwen3DecoderLayerImpl : public BaseLayer { void set_layer_id(int32_t layer_id) override { prefill_param_.layerId = layer_id; decode_graph_param_.layerId = layer_id; + decode_optimized_graph_param_.layerId = layer_id; decode_eager_param_.layerId = layer_id; } @@ -97,11 +98,14 @@ class NpuQwen3DecoderLayerImpl : public BaseLayer { atb_speed::Model::Node prefill_node_; atb_speed::Model::Node decode_graph_node_; + atb_speed::Model::Node decode_optimized_graph_node_; atb_speed::Model::Node decode_eager_node_; std::string model_name_; atb_speed::qwen::QwenLayerParam prefill_param_; atb_speed::qwen::QwenLayerParam decode_graph_param_; + atb_speed::qwen::QwenLayerParam decode_optimized_graph_param_; atb_speed::qwen::QwenLayerParam decode_eager_param_; + bool enable_optimized_decode_graph_ = false; atb::Tensor internal_tensors_; atb::Tensor residual_tensors_; atb::Tensor placeholder_; From c178f17df126bfdda2c0be32687bff69c70e68db Mon Sep 17 00:00:00 2001 From: Maojun Xu Date: Thu, 20 Aug 2026 15:07:05 +0800 Subject: [PATCH 2/3] perf: update optimized decode matmul revision. --- third_party/xllm_ops | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/xllm_ops b/third_party/xllm_ops index a1a6de04a5..aafec86d71 160000 --- a/third_party/xllm_ops +++ b/third_party/xllm_ops @@ -1 +1 @@ -Subproject commit a1a6de04a5934ad2c79cc8e2c7d222e62cf21088 +Subproject commit aafec86d719bfc77cb7e19d87e83206aee8a9554 From 850b87745478591874affe85dfcbfd71c577ff79 Mon Sep 17 00:00:00 2001 From: Maojun Xu Date: Thu, 20 Aug 2026 17:11:15 +0800 Subject: [PATCH 3/3] perf: specialize quantized decode shapes --- third_party/xllm_ops | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/xllm_ops b/third_party/xllm_ops index aafec86d71..74686ce1b1 160000 --- a/third_party/xllm_ops +++ b/third_party/xllm_ops @@ -1 +1 @@ -Subproject commit aafec86d719bfc77cb7e19d87e83206aee8a9554 +Subproject commit 74686ce1b1cca84c3554490eef35b770836f5b2e