[codex] Fix Kunlun Qwen2.5 greedy inference

csrc/engine/rank_worker.cpp

@@ -1,11 +1,69 @@
 #include "rank_worker.hpp"
 #include "../models/model_factory.hpp"
 #include "infinicore/ops.hpp"
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+#include <cstring>
+#include <limits>
 #include <spdlog/spdlog.h>
 #include <stdexcept>
 
 namespace infinilm::engine {
 
+namespace {
+
+float f16_to_f32(uint16_t h) {
+    uint32_t sign = (h & 0x8000u) << 16;
+    uint32_t exp = (h & 0x7c00u) >> 10;
+    uint32_t mant = h & 0x03ffu;
+    uint32_t bits;
+
+    if (exp == 0) {
+        if (mant == 0) {
+            bits = sign;
+        } else {
+            exp = 1;
+            while ((mant & 0x0400u) == 0) {
+                mant <<= 1;
+                --exp;
+            }
+            mant &= 0x03ffu;
+            bits = sign | ((exp + 127 - 15) << 23) | (mant << 13);
+        }
+    } else if (exp == 0x1fu) {
+        bits = sign | 0x7f800000u | (mant << 13);
+    } else {
+        bits = sign | ((exp + 127 - 15) << 23) | (mant << 13);
+    }
+
+    float out;
+    std::memcpy(&out, &bits, sizeof(out));
+    return out;
+}
+
+float bf16_to_f32(uint16_t h) {
+    uint32_t bits = static_cast<uint32_t>(h) << 16;
+    float out;
+    std::memcpy(&out, &bits, sizeof(out));
+    return out;
+}
+
+float read_scalar_as_f32(const std::byte *data, infinicore::DataType dtype, size_t idx) {
+    switch (dtype) {
+    case infinicore::DataType::F32:
+        return reinterpret_cast<const float *>(data)[idx];
+    case infinicore::DataType::F16:
+        return f16_to_f32(reinterpret_cast<const uint16_t *>(data)[idx]);
+    case infinicore::DataType::BF16:
+        return bf16_to_f32(reinterpret_cast<const uint16_t *>(data)[idx]);
+    default:
+        return std::numeric_limits<float>::quiet_NaN();
+    }
+}
+
+} // namespace
+
 RankWorker::RankWorker(
     std::shared_ptr<infinilm::global_state::InfinilmConfig> infinilm_config,
     const distributed::RankInfo &rank_info,
@@ -426,13 +484,75 @@ void RankWorker::thread_loop() {
                             const auto &total_len{logits_shape[1]};
                             const auto &batch_size{logits_shape[0]};
 
-                            auto n_req = local_args.input_offsets.value()->size(0) - 1;
-                            int32_t *input_offsets = (int32_t *)local_args.input_offsets.value()->data();
+                            auto input_offsets_cpu = local_args.input_offsets.value()->to(infinicore::Device::cpu());
+                            auto n_req = input_offsets_cpu->size(0) - 1;
+                            int32_t *input_offsets = (int32_t *)input_offsets_cpu->data();
 
-                            auto output_ids{infinicore::Tensor::empty({n_req}, infinicore::DataType::I64, rank_info_.device)};
+                            bool use_cpu_argmax = top_k == 1 && rank_info_.device.getType() == infinicore::Device::Type::KUNLUN;
+                            auto output_ids{infinicore::Tensor::empty(
+                                {n_req},
+                                infinicore::DataType::I64,
+                                use_cpu_argmax ? infinicore::Device::cpu() : rank_info_.device)};
+                            auto *output_ids_cpu = use_cpu_argmax
+                                                     ? reinterpret_cast<int64_t *>(output_ids->data())
+                                                     : nullptr;
 
                             for (auto i{decltype(n_req)(0)}; i < n_req; ++i) {
                                 auto score{logits->view({batch_size * total_len, vocab_size})->narrow({{0, size_t(input_offsets[i + 1] - 1), 1}})->view({vocab_size})};
+                                if (use_cpu_argmax) {
+                                    auto score_cpu = score->contiguous()->to(infinicore::Device::cpu());
+                                    const auto *score_data = score_cpu->data();
+                                    float best_score = -std::numeric_limits<float>::infinity();
+                                    size_t best_token = 0;
+                                    for (size_t token_idx = 0; token_idx < vocab_size; ++token_idx) {
+                                        float token_score = read_scalar_as_f32(score_data, score_cpu->dtype(), token_idx);
+                                        if (token_score > best_score) {
+                                            best_score = token_score;
+                                            best_token = token_idx;
+                                        }
+                                    }
+
+                                    output_ids_cpu[i] = static_cast<int64_t>(best_token);
+
+                                    if (std::getenv("INFINILM_DEBUG_SAMPLE") != nullptr) {
+                                        spdlog::warn(
+                                            "sample req={} offset={} cpu_argmax token={} score={}",
+                                            i,
+                                            input_offsets[i + 1] - 1,
+                                            best_token,
+                                            best_score);
+                                    }
+                                    continue;
+                                }
+
+                                if (std::getenv("INFINILM_DEBUG_SAMPLE") != nullptr) {
+                                    auto score_cpu = score->contiguous()->to(infinicore::Device::cpu());
+                                    const auto *score_data = score_cpu->data();
+                                    std::vector<std::pair<float, size_t>> top_scores;
+                                    top_scores.reserve(vocab_size);
+                                    for (size_t token_idx = 0; token_idx < vocab_size; ++token_idx) {
+                                        top_scores.emplace_back(
+                                            read_scalar_as_f32(score_data, score_cpu->dtype(), token_idx),
+                                            token_idx);
+                                    }
+                                    size_t k = std::min<size_t>(5, top_scores.size());
+                                    std::partial_sort(
+                                        top_scores.begin(),
+                                        top_scores.begin() + k,
+                                        top_scores.end(),
+                                        [](const auto &a, const auto &b) {
+                                            return a.first > b.first;
+                                        });
+                                    for (size_t j = 0; j < k; ++j) {
+                                        spdlog::warn(
+                                            "sample req={} offset={} top{} token={} score={}",
+                                            i,
+                                            input_offsets[i + 1] - 1,
+                                            j,
+                                            top_scores[j].second,
+                                            top_scores[j].first);
+                                    }
+                                }
                                 auto out{output_ids->narrow({{0, i, 1}})->view({})};
                                 float random_val = std::uniform_real_distribution<float>(0, 1)(rng_);
                                 infinicore::op::random_sample_(

csrc/layers/attention/attention.cpp

@@ -27,11 +27,15 @@ Attention::Attention(std::shared_ptr<infinilm::config::ModelConfig> model_config
     num_key_value_heads_ = total_num_kv_heads < tp_size ? 1 : total_num_kv_heads / tp_size;
 
     auto quantization_method = model_config->get_quantization_method();
-    auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); };
-    qkv_proj_ = std::make_shared<layers::linear::QKVParallelLinear>(
-        hidden_size_, head_dim_, total_num_heads, total_num_kv_heads,
-        "q_proj", "k_proj", "v_proj", register_fn,
-        quantization_method, use_bias, dtype, device, rank_info);
+    q_proj_ = this->register_module<layers::linear::ColumnParallelLinear>(
+        "q_proj", hidden_size_, total_num_heads * head_dim_, quantization_method,
+        use_bias, dtype, device, tp_rank, tp_size);
+    k_proj_ = this->register_module<layers::linear::ColumnParallelLinear>(
+        "k_proj", hidden_size_, total_num_kv_heads * head_dim_, quantization_method,
+        use_bias, dtype, device, tp_rank, tp_size);
+    v_proj_ = this->register_module<layers::linear::ColumnParallelLinear>(
+        "v_proj", hidden_size_, total_num_kv_heads * head_dim_, quantization_method,
+        use_bias, dtype, device, tp_rank, tp_size);
     o_proj_ = this->register_module<layers::linear::RowParallelLinear>(
         "o_proj", total_num_heads * head_dim_, hidden_size_, quantization_method,
         use_output_bias, dtype, device, tp_rank, tp_size, rank_info.comm);
@@ -42,6 +46,7 @@ Attention::Attention(std::shared_ptr<infinilm::config::ModelConfig> model_config
     attn_ = std::make_shared<AttentionLayer>(num_attention_heads_, head_dim_, scaling, num_key_value_heads_, layer_idx_,
                                              kv_cache_k_scale_, kv_cache_v_scale_, attention_backend_);
 
+    auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); };
     init_kv_cache_quant_params(register_fn, device, kv_cache_k_scale_, kv_cache_v_scale_);
 }
 
@@ -62,7 +67,9 @@ infinicore::Tensor Attention::forward_static_(const infinicore::Tensor &position
     size_t seq_len = shape[1];
 
     // 1. Project Q, K, V
-    auto [q, k, v] = qkv_proj_->forward_split(hidden_states_mutable);
+    auto q = q_proj_->forward(hidden_states_mutable);
+    auto k = k_proj_->forward(hidden_states_mutable);
+    auto v = v_proj_->forward(hidden_states_mutable);
 
     // 2. Reshape for multi-head attention
     auto q_reshaped = q->view({batch_size, seq_len, num_attention_heads_, head_dim_});
@@ -82,8 +89,7 @@ infinicore::Tensor Attention::forward_static_(const infinicore::Tensor &position
     }
 
     // 4. Apply RoPE to QK
-    auto q_rope = infinicore::Tensor::empty({batch_size, num_attention_heads_, seq_len, head_dim_}, q_reshaped->dtype(), q_reshaped->device())->permute({0, 2, 1, 3});
-    rotary_emb_->forward(q_rope, q_reshaped, pos_ids_for_rope);
+    auto q_rope = rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true);
     rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true);
 
     // 5. Attn Backend calculate
@@ -106,7 +112,9 @@ infinicore::Tensor Attention::forward_paged_(const infinicore::Tensor &position_
     ASSERT_EQ(batch_size, 1);
 
     // 1. Project Q, K, V
-    auto [q, k, v] = qkv_proj_->forward_split(hidden_states_mutable);
+    auto q = q_proj_->forward(hidden_states_mutable);
+    auto k = k_proj_->forward(hidden_states_mutable);
+    auto v = v_proj_->forward(hidden_states_mutable);
 
     // 2. Reshape for multi-head attention
     auto q_reshaped = q->view({seq_len, num_attention_heads_, head_dim_});

csrc/layers/attention/attention.hpp

@@ -21,11 +21,15 @@ class Attention : public infinicore::nn::Module {
                                const infinicore::Tensor &hidden_states) const;
 
     void process_weights_after_loading() override {
-        qkv_proj_->process_weights_after_loading();
+        q_proj_->process_weights_after_loading();
+        k_proj_->process_weights_after_loading();
+        v_proj_->process_weights_after_loading();
     }
 
     void reset_runtime_state() const override {
-        qkv_proj_->reset_runtime_state();
+        q_proj_->reset_runtime_state();
+        k_proj_->reset_runtime_state();
+        v_proj_->reset_runtime_state();
     }
 
     size_t layer_idx() const { return layer_idx_; }
@@ -42,7 +46,9 @@ class Attention : public infinicore::nn::Module {
                                       const infinicore::Tensor &hidden_states) const;
 
 protected:
-    std::shared_ptr<infinilm::layers::linear::QKVParallelLinear> qkv_proj_;
+    std::shared_ptr<infinilm::layers::linear::ColumnParallelLinear> q_proj_;
+    std::shared_ptr<infinilm::layers::linear::ColumnParallelLinear> k_proj_;
+    std::shared_ptr<infinilm::layers::linear::ColumnParallelLinear> v_proj_;
     std::shared_ptr<infinilm::layers::linear::RowParallelLinear> o_proj_;
     std::shared_ptr<infinicore::nn::RoPE> rotary_emb_;
 

csrc/layers/attention/backends/static_attn.cpp

@@ -74,8 +74,14 @@ infinicore::Tensor StaticAttentionImpl::forward(const AttentionLayer &layer,
                 q_reshaped);
         }
 
-        k_total = k_total->narrow({{2, 0, total_seq_len}}); // [bs, n_kv_head, total_seq_len, head_dim]
-        v_total = v_total->narrow({{2, 0, total_seq_len}}); // [bs, n_kv_head, total_seq_len, head_dim]
+        const size_t past_seq_len = reinterpret_cast<int32_t *>(past_sequence_lengths.value()->to(infinicore::Device::cpu())->data())[0];
+        if (past_seq_len == 0 && total_seq_len == seq_len) {
+            k_total = k_permuted->contiguous();
+            v_total = v_permuted->contiguous();
+        } else {
+            k_total = k_total->narrow({{2, 0, total_seq_len}})->contiguous(); // [bs, n_kv_head, total_seq_len, head_dim]
+            v_total = v_total->narrow({{2, 0, total_seq_len}})->contiguous(); // [bs, n_kv_head, total_seq_len, head_dim]
+        }
 
         //  Compute attention
         size_t ngroup = num_heads_ / num_kv_heads_;

csrc/layers/mlp/mlp.cpp

@@ -17,10 +17,12 @@ MLP::MLP(std::shared_ptr<infinilm::config::ModelConfig> model_config,
     int tp_size = rank_info.tp_size;
 
     auto quantization_method = model_config->get_quantization_method();
-    auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); };
-    gate_up_proj_ = std::make_shared<layers::linear::GateUpParallelLinear>(
-        hidden_size_, intermediate_size_, "gate_proj", "up_proj", register_fn,
-        quantization_method, use_bias_, dtype, device, rank_info);
+    gate_proj_ = this->register_module<layers::linear::ColumnParallelLinear>(
+        "gate_proj", hidden_size_, intermediate_size_, quantization_method,
+        use_bias_, dtype, device, tp_rank, tp_size);
+    up_proj_ = this->register_module<layers::linear::ColumnParallelLinear>(
+        "up_proj", hidden_size_, intermediate_size_, quantization_method,
+        use_bias_, dtype, device, tp_rank, tp_size);
     down_proj_ = this->register_module<layers::linear::RowParallelLinear>(
         "down_proj", intermediate_size_, hidden_size_, quantization_method,
         use_bias_, dtype, device, tp_rank, tp_size, rank_info.comm);
@@ -29,7 +31,8 @@ MLP::MLP(std::shared_ptr<infinilm::config::ModelConfig> model_config,
 infinicore::Tensor MLP::forward(const infinicore::Tensor &hidden_states) const {
     // 1. Project to gate and up
     auto hidden_states_mutable = hidden_states;
-    auto [gate, up] = gate_up_proj_->forward_split(hidden_states_mutable);
+    auto gate = gate_proj_->forward(hidden_states_mutable);
+    auto up = up_proj_->forward(hidden_states_mutable);
     // 2. Apply SwiGLU: silu(gate) * up
     auto intermediate = infinicore::op::swiglu(up, gate);
     // 3. Project down

csrc/layers/mlp/mlp.hpp

@@ -37,19 +37,22 @@ class MLP : public infinicore::nn::Module {
     infinicore::Tensor forward(const infinicore::Tensor &hidden_states) const;
 
     void process_weights_after_loading() override {
-        gate_up_proj_->process_weights_after_loading();
+        gate_proj_->process_weights_after_loading();
+        up_proj_->process_weights_after_loading();
     }
 
     void reset_runtime_state() const override {
-        gate_up_proj_->reset_runtime_state();
+        gate_proj_->reset_runtime_state();
+        up_proj_->reset_runtime_state();
     }
 
     // Module information
     size_t hidden_size() const { return hidden_size_; }
     size_t intermediate_size() const { return intermediate_size_; }
 
 protected:
-    std::shared_ptr<infinilm::layers::linear::GateUpParallelLinear> gate_up_proj_;
+    std::shared_ptr<infinilm::layers::linear::ColumnParallelLinear> gate_proj_;
+    std::shared_ptr<infinilm::layers::linear::ColumnParallelLinear> up_proj_;
     std::shared_ptr<infinilm::layers::linear::RowParallelLinear> down_proj_;
 
     size_t hidden_size_;

csrc/models/videonsa/videonsa_attention.cpp

@@ -281,7 +281,9 @@ infinicore::Tensor VideoNSAAttention::forward(const infinicore::Tensor &position
     auto hidden_states_mutable = hidden_states;
     const size_t batch_size = hidden_states->size(0);
     const size_t seq_len = hidden_states->size(1);
-    auto [q, k, v] = qkv_proj_->forward_split(hidden_states_mutable);
+    auto q = q_proj_->forward(hidden_states_mutable);
+    auto k = k_proj_->forward(hidden_states_mutable);
+    auto v = v_proj_->forward(hidden_states_mutable);
 
     auto pos_shape = positions->shape();
     infinicore::Tensor pos_ids_for_rope = positions;

examples/test_infer.py

@@ -10,6 +10,7 @@ def test(
     model_path,
     max_new_tokens=100,
     device="cpu",
+    dtype="float16",
     tp=1,
     enable_paged_attn=False,
     enable_graph=False,
@@ -34,6 +35,7 @@ def test(
     model = LLM(
         model_path=model_path,
         device=device,
+        dtype=dtype,
         tensor_parallel_size=tp,
         cache_type="paged" if enable_paged_attn else "static",
         max_batch_size=len(prompts),
@@ -107,6 +109,7 @@ def test(
         model_path,
         max_new_tokens,
         device=device_str,
+        dtype=cfg.dtype,
         tp=tp,
         enable_paged_attn=enable_paged_attn,
         enable_graph=enable_graph,

python/infinilm/base_config.py

@@ -484,7 +484,7 @@ def get_device_str(self, device):
             "metax": "cuda",
             "moore": "musa",
             "iluvatar": "cuda",
-            "kunlun": "kunlun",
+            "kunlun": "cuda",
             "hygon": "cuda",
             "ali": "cuda",
         }

python/infinilm/infer_engine.py

@@ -63,6 +63,18 @@ def read_hf_config(model_path):
 
     return config_dict
 
+def read_hf_config_with_dtype(model_path, dtype_override=None):
+    config_dict = read_hf_config(model_path)
+    if dtype_override is not None:
+        config_dict["torch_dtype"] = dtype_override
+        config_dict["dtype"] = dtype_override
+        if "text_config" in config_dict:
+            text_config = dict(config_dict["text_config"])
+            text_config["torch_dtype"] = dtype_override
+            text_config["dtype"] = dtype_override
+            config_dict["text_config"] = text_config
+    return config_dict
+
 # config.json (required) defines model architecture, while generation_config.json
 # (optional) defines generation behavior. They are kept as separate readers
 # because: 1) config.json must exist and requires model_type validation,
@@ -99,10 +111,11 @@ def __init__(
         enable_graph_compiling=False,
         attention_backend="default",
         kv_cache_dtype=None,
+        dtype=None,
         use_mla=False,
         weight_load_mode="async",
     ):
-        self.hf_config = read_hf_config(model_path)
+        self.hf_config = read_hf_config_with_dtype(model_path, dtype)
         self.hf_generation_config = read_hf_generation_config(model_path)
 
         if device is None:

python/infinilm/llm/model_runner/model_runner.py

@@ -73,6 +73,7 @@ def __init__(self, config: EngineConfig):
             cache_config=cache_config,
             enable_graph_compiling=config.enable_graph,
             attention_backend=config.attn_backend,
+            dtype=config.dtype,
             use_mla=config.use_mla,
             weight_load_mode=config.weight_load_mode,
         )