Skip to content

perf: route low-token qwen3 decode through optimized nz quantmatmul. #2248

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
maojunx99 wants to merge 3 commits into
base: main
Choose a base branch
from
Draft

perf: route low-token qwen3 decode through optimized nz quantmatmul. #2248

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
2bce490
perf: enable optimized qwen3 nz decode routing.
maojunx99 Aug 17, 2026
c178f17
perf: update optimized decode matmul revision.
maojunx99 Aug 20, 2026
850b877
perf: specialize quantized decode shapes
maojunx99 Aug 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion third_party/xllm_atb_layers
Open in desktop
Submodule xllm_atb_layers updated from 92171d to 760399
2 changes: 1 addition & 1 deletion third_party/xllm_ops
Open in desktop
Submodule xllm_ops updated from d4a441 to 74686c
84 changes: 81 additions & 3 deletions xllm/core/layers/npu/npu_qwen3_decoder_layer_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand All @@ -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,
Expand Down Expand Up @@ -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 =
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -227,7 +280,15 @@ int64_t NpuQwen3DecoderLayerImpl::init_layer() {
if (quantize_type_ == "w8a8") {
Qwen3DecoderLoader* qwen3_loader =
dynamic_cast<Qwen3DecoderLoader*>(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<int>(LinearTypeV2::W8A8);
Expand All @@ -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()) {
Expand All @@ -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_));

Expand Down Expand Up @@ -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,
Expand Down
4 changes: 4 additions & 0 deletions xllm/core/layers/npu/npu_qwen3_decoder_layer_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand Down Expand Up @@ -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_;
Expand Down