From deceab5025272ac7d1d389f48a44e507c31e77a8 Mon Sep 17 00:00:00 2001 From: Codex Date: Mon, 6 Jul 2026 02:12:13 +0000 Subject: [PATCH] Fix Kunlun Qwen2.5 greedy inference --- csrc/engine/rank_worker.cpp | 126 +++++++++++++++++- csrc/layers/attention/attention.cpp | 26 ++-- csrc/layers/attention/attention.hpp | 12 +- .../layers/attention/backends/static_attn.cpp | 10 +- csrc/layers/mlp/mlp.cpp | 13 +- csrc/layers/mlp/mlp.hpp | 9 +- csrc/models/videonsa/videonsa_attention.cpp | 4 +- examples/test_infer.py | 3 + python/infinilm/base_config.py | 2 +- python/infinilm/infer_engine.py | 15 ++- .../infinilm/llm/model_runner/model_runner.py | 1 + python/infinilm/modeling_utils.py | 7 +- .../processors/basic_llm_processor.py | 4 +- 13 files changed, 201 insertions(+), 31 deletions(-) diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index 52bc237d7..6c4b76361 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -1,11 +1,69 @@ #include "rank_worker.hpp" #include "../models/model_factory.hpp" #include "infinicore/ops.hpp" +#include +#include +#include +#include +#include #include #include 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(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(data)[idx]; + case infinicore::DataType::F16: + return f16_to_f32(reinterpret_cast(data)[idx]); + case infinicore::DataType::BF16: + return bf16_to_f32(reinterpret_cast(data)[idx]); + default: + return std::numeric_limits::quiet_NaN(); + } +} + +} // namespace + RankWorker::RankWorker( std::shared_ptr 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(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::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(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> 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(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(0, 1)(rng_); infinicore::op::random_sample_( diff --git a/csrc/layers/attention/attention.cpp b/csrc/layers/attention/attention.cpp index 1b87f6fbc..be8deaa83 100644 --- a/csrc/layers/attention/attention.cpp +++ b/csrc/layers/attention/attention.cpp @@ -27,11 +27,15 @@ Attention::Attention(std::shared_ptr 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( - 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( + "q_proj", hidden_size_, total_num_heads * head_dim_, quantization_method, + use_bias, dtype, device, tp_rank, tp_size); + k_proj_ = this->register_module( + "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( + "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( "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 model_config attn_ = std::make_shared(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_}); diff --git a/csrc/layers/attention/attention.hpp b/csrc/layers/attention/attention.hpp index 671b90f5c..95a8386fa 100644 --- a/csrc/layers/attention/attention.hpp +++ b/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 qkv_proj_; + std::shared_ptr q_proj_; + std::shared_ptr k_proj_; + std::shared_ptr v_proj_; std::shared_ptr o_proj_; std::shared_ptr rotary_emb_; diff --git a/csrc/layers/attention/backends/static_attn.cpp b/csrc/layers/attention/backends/static_attn.cpp index b71f0b52e..3dac78eaa 100644 --- a/csrc/layers/attention/backends/static_attn.cpp +++ b/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(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_; diff --git a/csrc/layers/mlp/mlp.cpp b/csrc/layers/mlp/mlp.cpp index f7604c505..7559bbbc7 100644 --- a/csrc/layers/mlp/mlp.cpp +++ b/csrc/layers/mlp/mlp.cpp @@ -17,10 +17,12 @@ MLP::MLP(std::shared_ptr 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( - hidden_size_, intermediate_size_, "gate_proj", "up_proj", register_fn, - quantization_method, use_bias_, dtype, device, rank_info); + gate_proj_ = this->register_module( + "gate_proj", hidden_size_, intermediate_size_, quantization_method, + use_bias_, dtype, device, tp_rank, tp_size); + up_proj_ = this->register_module( + "up_proj", hidden_size_, intermediate_size_, quantization_method, + use_bias_, dtype, device, tp_rank, tp_size); down_proj_ = this->register_module( "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 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 diff --git a/csrc/layers/mlp/mlp.hpp b/csrc/layers/mlp/mlp.hpp index abd81bd88..9f396bfdf 100644 --- a/csrc/layers/mlp/mlp.hpp +++ b/csrc/layers/mlp/mlp.hpp @@ -37,11 +37,13 @@ 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 @@ -49,7 +51,8 @@ class MLP : public infinicore::nn::Module { size_t intermediate_size() const { return intermediate_size_; } protected: - std::shared_ptr gate_up_proj_; + std::shared_ptr gate_proj_; + std::shared_ptr up_proj_; std::shared_ptr down_proj_; size_t hidden_size_; diff --git a/csrc/models/videonsa/videonsa_attention.cpp b/csrc/models/videonsa/videonsa_attention.cpp index 828d10860..e6f97ea5d 100644 --- a/csrc/models/videonsa/videonsa_attention.cpp +++ b/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; diff --git a/examples/test_infer.py b/examples/test_infer.py index c8c74b4fc..57479eb38 100644 --- a/examples/test_infer.py +++ b/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, diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 8166f8336..a231c5bdc 100644 --- a/python/infinilm/base_config.py +++ b/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", } diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index b45637f75..73729a3bc 100644 --- a/python/infinilm/infer_engine.py +++ b/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: diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index daf60e143..d0be4054d 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/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, ) diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index 33b4dc3c0..875f00a63 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -106,7 +106,12 @@ def load_state_dict( ) for k in f.keys(): - state_dict[k] = f.get_tensor(k).to(device=device) + tensor = f.get_tensor(k) + if tensor.is_floating_point(): + tensor = tensor.to(device=device, dtype=dtype) + else: + tensor = tensor.to(device=device) + state_dict[k] = tensor return state_dict diff --git a/python/infinilm/processors/basic_llm_processor.py b/python/infinilm/processors/basic_llm_processor.py index 6948aa41c..9c9db0d83 100644 --- a/python/infinilm/processors/basic_llm_processor.py +++ b/python/infinilm/processors/basic_llm_processor.py @@ -138,7 +138,7 @@ def _build_model_input_from_static_scheduler_output( return { "input_ids": infinicore.from_list(input_ids, dtype=infinicore.int64), - "position_ids": infinicore.from_list(position_ids, dtype=infinicore.int64), + "position_ids": infinicore.from_list(position_ids, dtype=infinicore.int32), "past_kv_lengths": infinicore.from_list( [past_kv_len], dtype=infinicore.int32 ), @@ -248,7 +248,7 @@ def _build_model_input_from_batch_scheduler_output( return { "input_ids": infinicore.from_list([tokens], dtype=infinicore.int64), - "position_ids": infinicore.from_list(position_ids, dtype=infinicore.int64), + "position_ids": infinicore.from_list(position_ids, dtype=infinicore.int32), "past_kv_lengths": infinicore.from_list( cached_lens, dtype=infinicore.int32 ),