perf: add NineToothed fused QK norm

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<double>("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<std::remove_reference_t<decltype(*q_norm_)>>(runtime, head_dim_, rms_norm_eps, dtype));
         k_norm_ = this->bind_child("k_norm", std::make_shared<std::remove_reference_t<decltype(*k_norm_)>>(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<infinilm::runtime::Tensor, infinilm::runtime::Tensor, infinilm::runtime::Tensor>
+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<infini::ops::FusedQkNorm>(
+        runtime, infini::ops::FusedQkNorm::DefaultConfig(runtime.device().type()),
+        packed_qkv, static_cast<int64_t>(num_attention_heads_),
+        static_cast<int64_t>(num_key_value_heads_),
+        static_cast<int64_t>(num_key_value_heads_),
+        static_cast<int64_t>(head_dim_),
+        static_cast<float>(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

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<infinilm::runtime::Tensor, infinilm::runtime::Tensor, infinilm::runtime::Tensor>
+    forward_split_fused_qk_(infinilm::runtime::Runtime &runtime,
+                            const infinilm::runtime::Tensor &qkv,
+                            size_t token_count) const;
+
 protected:
     std::shared_ptr<infinilm::layers::linear::QKVParallelLinear> qkv_proj_;
     std::shared_ptr<infinilm::layers::linear::RowParallelLinear> 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_;