From c0e63596a43f0adcabb140d820b7b7b6da0c14af Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Thu, 5 Nov 2026 16:42:27 +0800 Subject: [PATCH] perf: add NineToothed fused QK norm --- csrc/models/qwen3/qwen3_attention.cpp | 54 +++++++++++++++++++++++---- csrc/models/qwen3/qwen3_attention.hpp | 7 ++++ third_party/InfiniCore | 2 +- 3 files changed, 55 insertions(+), 8 deletions(-) diff --git a/csrc/models/qwen3/qwen3_attention.cpp b/csrc/models/qwen3/qwen3_attention.cpp index c276d3161..8fb49dbd8 100644 --- a/csrc/models/qwen3/qwen3_attention.cpp +++ b/csrc/models/qwen3/qwen3_attention.cpp @@ -1,4 +1,5 @@ #include "qwen3_attention.hpp" +#include "../../execution/submit.hpp" #include "../../global_state/global_state.hpp" #include "../../layers/attention/attention.hpp" #include "../../utils.hpp" @@ -25,6 +26,7 @@ Qwen3Attention::Qwen3Attention(infinilm::runtime::Runtime &runtime, std::shared_ double rms_norm_eps = model_config->get("rms_norm_eps"); attention_backend_ = attention_backend; + device_ = runtime.device(); int tp_rank = rank_info.tp_rank; int tp_size = rank_info.tp_size; num_attention_heads_ = total_num_heads / tp_size; @@ -49,6 +51,7 @@ Qwen3Attention::Qwen3Attention(infinilm::runtime::Runtime &runtime, std::shared_ q_norm_ = this->bind_child("q_norm", std::make_shared>(runtime, head_dim_, rms_norm_eps, dtype)); k_norm_ = this->bind_child("k_norm", std::make_shared>(runtime, head_dim_, rms_norm_eps, dtype)); + use_fused_qk_norm_ = device_.type() == infinilm::runtime::Device::Type::kNvidia && head_dim_ == 128 && dtype == infinilm::runtime::DataType::kBFloat16 && !infini::ops::FusedQkNorm::active_implementation_indices(device_.type()).empty(); // NONE KV scales remain the original null Parameters; no Device allocation. if (runtime.device().type() != infinilm::runtime::Device::Type::kCpu) { @@ -88,10 +91,16 @@ infinilm::runtime::Tensor Qwen3Attention::forward_static_(const infinilm::runtim size_t seq_len = shape[1]; // 1. Project Q, K, V - auto [q, k, v] = qkv_proj_->forward_split(*forward_context.runtime, hidden_states_mutable); - - q = q_norm_->forward(*forward_context.runtime, q->view({batch_size * seq_len, num_attention_heads_, head_dim_})); - k = k_norm_->forward(*forward_context.runtime, k->view({batch_size * seq_len, num_key_value_heads_, head_dim_})); + auto [q, k, v] = use_fused_qk_norm_ + ? forward_split_fused_qk_(*forward_context.runtime, + qkv_proj_->forward(*forward_context.runtime, hidden_states_mutable), + batch_size * seq_len) + : qkv_proj_->forward_split(*forward_context.runtime, hidden_states_mutable); + + if (!use_fused_qk_norm_) { + q = q_norm_->forward(*forward_context.runtime, q->view({batch_size * seq_len, num_attention_heads_, head_dim_})); + k = k_norm_->forward(*forward_context.runtime, k->view({batch_size * seq_len, num_key_value_heads_, head_dim_})); + } // 2. Reshape for multi-head attention auto q_reshaped = q->view({batch_size, seq_len, num_attention_heads_, head_dim_}); @@ -139,14 +148,20 @@ infinilm::runtime::Tensor Qwen3Attention::forward_paged_(const infinilm::runtime ASSERT_EQ(batch_size, 1); // 1. Project Q, K, V - auto [q, k, v] = qkv_proj_->forward_split(*forward_context.runtime, hidden_states_mutable); + auto [q, k, v] = use_fused_qk_norm_ + ? forward_split_fused_qk_(*forward_context.runtime, + qkv_proj_->forward(*forward_context.runtime, hidden_states_mutable), + seq_len) + : qkv_proj_->forward_split(*forward_context.runtime, hidden_states_mutable); // 2. Reshape for multi-head attention auto q_reshaped = q->view({seq_len, num_attention_heads_, head_dim_}); auto k_reshaped = k->view({seq_len, num_key_value_heads_, head_dim_}); auto v_reshaped = v->view({seq_len, num_key_value_heads_, head_dim_}); - q_reshaped = q_norm_->forward(*forward_context.runtime, q_reshaped); - k_reshaped = k_norm_->forward(*forward_context.runtime, k_reshaped); + if (!use_fused_qk_norm_) { + q_reshaped = q_norm_->forward(*forward_context.runtime, q_reshaped); + k_reshaped = k_norm_->forward(*forward_context.runtime, k_reshaped); + } // 3. Prepare position_ids for RoPE auto pos_shape = position_ids->shape(); @@ -171,4 +186,29 @@ infinilm::runtime::Tensor Qwen3Attention::forward_paged_(const infinilm::runtime auto output = o_proj_->forward(*forward_context.runtime, attn_output); return output; } + +std::tuple +Qwen3Attention::forward_split_fused_qk_(infinilm::runtime::Runtime &runtime, + const infinilm::runtime::Tensor &qkv, + size_t token_count) const { + const size_t q_width = num_attention_heads_ * head_dim_; + const size_t k_width = num_key_value_heads_ * head_dim_; + const size_t total_width = qkv->shape()[2]; + + auto packed_qkv = qkv->view({token_count, total_width}); + infinilm::execution::submit( + runtime, infini::ops::FusedQkNorm::DefaultConfig(runtime.device().type()), + packed_qkv, static_cast(num_attention_heads_), + static_cast(num_key_value_heads_), + static_cast(num_key_value_heads_), + static_cast(head_dim_), + static_cast(q_norm_->eps()), q_norm_->weight(), k_norm_->weight()); + + auto q = qkv->narrow({{2, 0, q_width}}) + ->view({token_count, num_attention_heads_, head_dim_}); + auto k = qkv->narrow({{2, q_width, k_width}}) + ->view({token_count, num_key_value_heads_, head_dim_}); + auto v = qkv->narrow({{2, q_width + k_width, qkv->shape()[2] - q_width - k_width}}); + return {q, k, v}; +} } // namespace infinilm::models::qwen3 diff --git a/csrc/models/qwen3/qwen3_attention.hpp b/csrc/models/qwen3/qwen3_attention.hpp index 7fd682f63..91ce28ce8 100644 --- a/csrc/models/qwen3/qwen3_attention.hpp +++ b/csrc/models/qwen3/qwen3_attention.hpp @@ -40,6 +40,11 @@ class Qwen3Attention : public infinilm::weights::ModelWeights { const infinilm::runtime::Tensor &hidden_states, infinilm::global_state::ForwardContext &forward_context) const; + std::tuple + forward_split_fused_qk_(infinilm::runtime::Runtime &runtime, + const infinilm::runtime::Tensor &qkv, + size_t token_count) const; + protected: std::shared_ptr qkv_proj_; std::shared_ptr o_proj_; @@ -54,6 +59,8 @@ class Qwen3Attention : public infinilm::weights::ModelWeights { size_t num_key_value_heads_; size_t hidden_size_; size_t head_dim_; + infinilm::runtime::Device device_; + bool use_fused_qk_norm_{false}; // For off-line kv cache quantization infinilm::weights::Weight kv_cache_k_scale_; diff --git a/third_party/InfiniCore b/third_party/InfiniCore index a6ba44fa5..82095baef 160000 --- a/third_party/InfiniCore +++ b/third_party/InfiniCore @@ -1 +1 @@ -Subproject commit a6ba44fa568244c5dcc0f12b799f609483c06205 +Subproject commit 82095baef0369797132898a8d132ca13ae60807b